CIS6270 / lecture_2 /velocity_fields.py

Upload 87 files

0ba9d09 verified about 1 month ago

14.5 kB

"""Velocity fields, flow maps and the continuity equation.
A velocity field assigns one vector to every state and every time, and we trace a trajectory by
following those vectors from a starting point. The trajectories from every starting point together
form the flow map, and the continuity equation relates the motion of the density to the same field.
Every field here is analytic, so every number printed has a closed form beside it.
Chapter 2, Sections 2.2 to 2.8 of the notes. Runs on the chapter's analytic fields in one and two
dimensions. Prints the trajectories, the Euler error at one, two and four steps and the gap closing
like 1/N, the pushforward of a Gaussian prior, the flux balance on [-1, 1], both sides of the
continuity equation at x = 2, and the log density along one trajectory.
    python lecture_2/velocity_fields.py                 # the full report
    python lecture_2/velocity_fields.py --steps 20      # a shorter Euler refinement ladder
"""

# %% 01. Imports and configuration
import argparse
import math

import torch

from cis6270.runner import print_report, save_run, set_seed

torch.set_default_dtype(torch.float64)  # Double precision, because the Euler ladder prints six decimal places.

LOG2 = math.log(2.0)  # The rate a of the expansion field v_t(x) = a x, so that e^{at} = 2^t.
GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.
START_POINTS = torch.tensor([-1.0, 0.0, 1.0])  # [3] the initial states followed through each field.
EULER_STEPS = 64  # Largest step count in the refinement ladder, halved down to one step.
NUM_SAMPLES = 4096  # Draws from the prior p_0 used for the pushforward moments.


# %% 02. Three velocity fields and their trajectories
#
# Section 2.2 and Section 2.3. Each field is a rule at a single instant, written as a function of
# the state x and the time t, and each has an exact solution to compare a solver against.
def constant_field(x, t):
    """v_t(x) = 2, the field that translates every state at the same speed."""
    return torch.full_like(x, 2.0)  # Same shape as x; the value depends on neither x nor t.


def linear_field(x, t, a=LOG2):
    """v_t(x) = a x. It expands away from the origin for a > 0 and contracts toward it for a < 0."""
    return a * x  # Speed grows with distance from the origin and vanishes at x = 0.


def rotation_field(x, t):
    """v(x_1, x_2) = (-x_2, x_1), everywhere perpendicular to the position vector."""
    return torch.stack([-x[..., 1], x[..., 0]], dim=-1)  # [..., 2] same shape as x.


def exact_linear(x_0, t, a=LOG2):
    """x_t = x_0 e^{at}, the solution of dx_t/dt = a x_t obtained by separating variables."""
    return x_0 * torch.exp(a * t)  # Broadcasts a grid of times against a set of starting points.


def trajectories():
    """The chapter's worked trajectories for translation, expansion, contraction and rotation."""
    speed = constant_field(torch.tensor(1.0), torch.tensor(0.0))  # The value 2, at any state.
    translated = 1.0 + speed * GRID  # [5] the integral form x_t = x_0 + 2t from x_0 = 1.
    expanded = exact_linear(START_POINTS, torch.tensor(1.0))  # [3] states at t = 1, each doubled.
    contracted = exact_linear(START_POINTS, torch.tensor(1.0), a=-LOG2)  # [3] states each halved.
    turns = torch.tensor([math.pi / 2, math.pi])  # [2] the quarter turn and the half turn.
    circle = torch.stack([torch.cos(turns), torch.sin(turns)], dim=-1)  # [2, 2] x_t = (cos t, sin t).
    # The squared norm is conserved because the field is orthogonal to the position: 2 x . v_t(x) = 0.
    off_axis = torch.tensor([0.3, 0.7])  # [2] a state away from the axes, to test the conservation.
    norm_rate = 2.0 * (off_axis * rotation_field(off_axis, torch.tensor(0.0))).sum()  # d||x||^2/dt.
    return {
        "constant_field_trajectory": [round(float(v), 4) for v in translated],
        "expansion_x_half_from_two": float(exact_linear(torch.tensor(2.0), torch.tensor(0.5))),
        "expansion_endpoints": [round(float(v), 4) + 0.0 for v in expanded],
        "expansion_outer_spacing": float(expanded[2] - expanded[0]),
        "contraction_endpoints": [round(float(v), 4) + 0.0 for v in contracted],
        "contraction_outer_spacing": float(contracted[2] - contracted[0]),
        "rotation_quarter_turn": [round(float(v), 4) + 0.0 for v in circle[0]],
        "rotation_half_turn": [round(float(v), 4) + 0.0 for v in circle[1]],
        "rotation_radius": float(circle.norm(dim=-1).mean()),
        "rotation_squared_norm_rate": float(norm_rate),
    }


# %% 03. The flow map
#
# Section 2.4. We solve the ODE from every starting point at once. That gives the flow map phi_t.
# At t = 0 that map is the identity, and for the expansion field it multiplies the whole line by
# 2^t.
def flow_map(x_0, t, a=LOG2):
    """phi_t(x_0) = x_0 e^{at}, the accumulated transformation after following v_t for a time t."""
    return exact_linear(x_0, t, a)  # [N] the state each starting point has reached by time t.


def flow_map_report():
    """The identity at t = 0, the Jacobian determinant at t = 1, and noncrossing trajectories."""
    probe = torch.linspace(-3.0, 3.0, 101)  # [101] states spread over the line.
    identity_error = (flow_map(probe, torch.tensor(0.0)) - probe).abs().max()  # phi_0 = id.
    x_0 = torch.tensor(1.3, requires_grad=True)  # A state at which to differentiate phi_1.
    jacobian = torch.autograd.grad(flow_map(x_0, torch.tensor(1.0)), x_0)[0]  # d phi_1 / d x_0 = 2.
    # Proposition 2.1: distinct trajectories of a Lipschitz field never share a state at a shared time.
    gaps = (flow_map(START_POINTS[2], GRID) - flow_map(START_POINTS[0], GRID)).abs()  # [5].
    return {
        "flow_map_identity_error": float(identity_error),
        "flow_map_jacobian_det": float(jacobian),
        "trajectory_gap_minimum": float(gaps.min()),
    }


# %% 04. Euler's method
#
# Section 2.5. Within one step the velocity stays fixed at its value at the start of the step. On
# dx_t/dt = x_t the velocity grows with the state, so every step undershoots.
def euler(field, x_0, steps):
    """hat x_{k+1} = hat x_k + Delta t v_{t_k}(hat x_k) on a uniform grid of ``steps`` intervals."""
    x_hat = x_0.clone()  # The running approximation, started from the exact initial state.
    dt = 1.0 / steps  # Delta t, the width of one interval of the grid.
    for k in range(steps):  # One velocity evaluation per step.
        x_hat = x_hat + dt * field(x_hat, torch.tensor(k * dt))  # The Euler update at grid time t_k.
    return x_hat


def euler_ladder(largest):
    """Euler on dx_t/dt = x_t from x_0 = 1, at every power of two up to ``largest`` steps."""
    exact = math.e  # The exact solution x_1 = e^1, the chapter's 2.7183.
    counts, landings, errors = [], [], []
    steps = 1
    largest = max(largest, 4)  # The chapter's three rungs, at one, two and four steps, always run.
    while steps <= largest:  # 1, 2, 4, 8, ... up to the requested largest step count.
        landing = float(euler(lambda x, t: x, torch.tensor(1.0), steps))  # The field is v_t(x) = x.
        counts.append(steps)
        landings.append(round(landing, 6))
        errors.append(round(exact - landing, 6))
        steps *= 2
    return {
        "euler_step_counts": counts,
        "euler_landings": landings,
        "euler_errors": errors,
        # A first-order method has error proportional to 1/N, so the error times N settles at e/2.
        "euler_error_times_steps": [round(e * n, 4) for e, n in zip(errors, counts)],
        "euler_error_halving_ratios": [round(errors[i] / errors[i + 1], 4)
                                       for i in range(len(errors) - 1)],
        "euler_relative_error_four_steps": float((exact - landings[2]) / exact),
    }


# %% 05. A random starting point
#
# Section 2.6. The initial state becomes a random variable. Each draw follows the same field, so
# the law at time t is the pushforward p_t = (phi_t)_# p_0.
def pushforward_moments(num_samples, seed):
    """Moments of phi_1(X_0) for X_0 ~ N(0, 1) under the expansion field, against N(0, 4)."""
    set_seed(seed)  # Same seed, same draw, same report.
    x_0 = torch.randn(num_samples)  # [N] draws from the prior p_0 = N(0, 1).
    x_1 = flow_map(x_0, torch.tensor(1.0))  # [N] states at t = 1, every one of them doubled.
    return {
        "pushforward_mean": float(x_1.mean()),  # Exact value 0, since the flow fixes the origin.
        "pushforward_variance": float(x_1.var(unbiased=False)),  # Scaling by 2 gives variance 4.
        "pushforward_exact_variance": 4.0,  # Scaling a centered law by c multiplies variance by c^2.
    }


# %% 06. Probability flux and the continuity equation
#
# Section 2.7. The flux j_t(x) = p_t(x) v_t(x) measures probability transported per unit time.
# Once an interval shrinks to a point, the balance over it becomes the continuity equation.
def spreading_density(x, t):
    """p_t(x) for X_t ~ N(0, 4^t), the law the expansion field produces from a standard normal."""
    std = torch.pow(torch.tensor(2.0), t)  # 2^t, the standard deviation at time t.
    return torch.exp(-x ** 2 / (2.0 * std ** 2)) / (std * math.sqrt(2.0 * math.pi))


def flux_balance():
    """The two endpoint fluxes on [-1, 1] at t = 0, and the rate of change of the enclosed mass."""
    t_0 = torch.tensor(0.0)
    endpoints = torch.tensor([-1.0, 1.0])  # [2] the interval's endpoints a and b.
    density = spreading_density(endpoints, t_0)  # [2] p_0(-1) and p_0(1), both 0.2420.
    flux = density * linear_field(endpoints, t_0)  # [2] j_0(x) = p_0(x) v_0(x).
    # The same rate read off the distribution: M_t = erf(2^{-t} / sqrt 2), differentiated in t.
    t = torch.tensor(0.0, requires_grad=True)
    mass = torch.erf(torch.pow(torch.tensor(2.0), -t) / math.sqrt(2.0))  # P(|X_t| <= 1).
    mass_rate = torch.autograd.grad(mass, t)[0]  # d M_t / dt at t = 0.
    return {
        "density_at_one": float(density[1]),
        "flux_at_minus_one": float(flux[0]),
        "flux_at_one": float(flux[1]),
        "flux_balance": float(flux[0] - flux[1]),  # j_0(a) - j_0(b), the net rate into [-1, 1].
        "enclosed_mass": float(mass.detach()),
        "enclosed_mass_rate": float(mass_rate),
    }


def continuity_check():
    """Both sides of the continuity equation at x = 2 and t = 0, in the expanded product-rule form."""
    x = torch.tensor(2.0, requires_grad=True)  # The position held fixed while the density changes.
    t = torch.tensor(0.0, requires_grad=True)  # The time at which both sides are evaluated.
    p = spreading_density(x, t)  # p_0(2), the chapter's 0.05399.
    dp_dt = torch.autograd.grad(p, t, create_graph=True)[0]  # Local time derivative of the density.
    dp_dx = torch.autograd.grad(p, x, create_graph=True)[0]  # Spatial gradient of the density.
    divergence = torch.autograd.grad(linear_field(x, t), x)[0]  # div v = log 2 in one dimension.
    transport = linear_field(x, t) * dp_dx  # v_t(x) . grad p_t(x), the transport term.
    volume = p * divergence  # p_t(x) div v_t(x), the local volume-change term.
    return {
        "density_at_two": float(p.detach()),
        "density_gradient_at_two": float(dp_dx.detach()),
        "velocity_divergence": float(divergence),
        "continuity_time_derivative": float(dp_dt.detach()),
        "continuity_minus_flux_divergence": float(-(transport + volume).detach()),
    }


# %% 07. The density along a moving trajectory
#
# Section 2.8. An observer riding with the sample sees no transport term, so the log density
# changes at the rate -div v_t(x_t). Over a unit of time that rate agrees with the Jacobian of the
# flow map.
def log_density_along_trajectory(x_start=0.7):
    """The instantaneous change-of-variables equation on the expansion field, both ways."""
    x_0 = torch.tensor(x_start)  # One starting point, carried from t = 0 to t = 1.
    x_1 = flow_map(x_0, torch.tensor(1.0))  # The state the trajectory reaches at t = 1.
    log_p0 = torch.log(spreading_density(x_0, torch.tensor(0.0)))  # log p_0(x_0).
    log_p1 = torch.log(spreading_density(x_1, torch.tensor(1.0)))  # log p_1(x_1).
    origin = torch.tensor(0.0)  # The fixed point of the flow, where phi_1(0) = 0.
    prior_at_origin = float(spreading_density(origin, torch.tensor(0.0)))  # p_0(0) = 0.3989.
    final_at_origin = float(spreading_density(origin, torch.tensor(1.0)))  # p_1(0) = 0.1995.
    return {
        "log_density_rate": -LOG2,  # d log p_t(x_t) / dt = -div v_t(x_t), a constant here.
        "log_density_change": float(log_p1 - log_p0),  # Measured from the two densities.
        "log_density_accumulated": -LOG2 * 1.0,  # The integral of -div v_s over s from 0 to 1.
        "prior_density_at_origin": prior_at_origin,
        "final_density_at_origin": final_at_origin,
        "density_ratio_at_origin": prior_at_origin / final_at_origin,  # |det J_{phi_1}| = 2.
    }


# %% 08. The report
def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--steps", type=int, default=EULER_STEPS, help="largest Euler step count")
    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="draws from the prior")
    parser.add_argument("--seed", type=int, default=0, help="random seed")
    parser.add_argument("--out", type=str, default="", help="directory for the run's report")
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    args = parser.parse_args()

    report = {}
    report.update(trajectories())  # Sections 2.2 and 2.3.
    report.update(flow_map_report())  # Section 2.4.
    report.update(euler_ladder(args.steps))  # Section 2.5.
    report.update(pushforward_moments(args.samples, args.seed))  # Section 2.6.
    report.update(flux_balance())  # Section 2.7.
    report.update(continuity_check())  # Section 2.7.
    report.update(log_density_along_trajectory())  # Section 2.8.

    if not args.quiet:  # The Euler ladder, rung by rung, as the chapter develops it.
        for n, landing, error in zip(report["euler_step_counts"], report["euler_landings"],
                                     report["euler_errors"]):
            print(f"  euler steps {n:3d}  landing {landing:.6f}  error {error:.6f}"
                  f"  error x steps {error * n:.4f}")

    print_report("velocity fields, flow maps and the continuity equation", report)
    save_run(args.out, report=report)  # Writes report.json when --out names a directory.


# %% 09. Run the script
if __name__ == "__main__":
    main()

	"""Velocity fields, flow maps and the continuity equation.

	A velocity field assigns one vector to every state and every time, and we trace a trajectory by
	following those vectors from a starting point. The trajectories from every starting point together
	form the flow map, and the continuity equation relates the motion of the density to the same field.
	Every field here is analytic, so every number printed has a closed form beside it.

	Chapter 2, Sections 2.2 to 2.8 of the notes. Runs on the chapter's analytic fields in one and two
	dimensions. Prints the trajectories, the Euler error at one, two and four steps and the gap closing
	like 1/N, the pushforward of a Gaussian prior, the flux balance on [-1, 1], both sides of the
	continuity equation at x = 2, and the log density along one trajectory.

	python lecture_2/velocity_fields.py # the full report
	python lecture_2/velocity_fields.py --steps 20 # a shorter Euler refinement ladder
	"""

	# %% 01. Imports and configuration
	import argparse
	import math

	import torch

	from cis6270.runner import print_report, save_run, set_seed

	torch.set_default_dtype(torch.float64) # Double precision, because the Euler ladder prints six decimal places.

	LOG2 = math.log(2.0) # The rate a of the expansion field v_t(x) = a x, so that e^{at} = 2^t.
	GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0]) # [5] the chapter's grid of times.
	START_POINTS = torch.tensor([-1.0, 0.0, 1.0]) # [3] the initial states followed through each field.
	EULER_STEPS = 64 # Largest step count in the refinement ladder, halved down to one step.
	NUM_SAMPLES = 4096 # Draws from the prior p_0 used for the pushforward moments.


	# %% 02. Three velocity fields and their trajectories
	#
	# Section 2.2 and Section 2.3. Each field is a rule at a single instant, written as a function of
	# the state x and the time t, and each has an exact solution to compare a solver against.
	def constant_field(x, t):
	"""v_t(x) = 2, the field that translates every state at the same speed."""
	return torch.full_like(x, 2.0) # Same shape as x; the value depends on neither x nor t.


	def linear_field(x, t, a=LOG2):
	"""v_t(x) = a x. It expands away from the origin for a > 0 and contracts toward it for a < 0."""
	return a * x # Speed grows with distance from the origin and vanishes at x = 0.


	def rotation_field(x, t):
	"""v(x_1, x_2) = (-x_2, x_1), everywhere perpendicular to the position vector."""
	return torch.stack([-x[..., 1], x[..., 0]], dim=-1) # [..., 2] same shape as x.


	def exact_linear(x_0, t, a=LOG2):
	"""x_t = x_0 e^{at}, the solution of dx_t/dt = a x_t obtained by separating variables."""
	return x_0 * torch.exp(a * t) # Broadcasts a grid of times against a set of starting points.


	def trajectories():
	"""The chapter's worked trajectories for translation, expansion, contraction and rotation."""
	speed = constant_field(torch.tensor(1.0), torch.tensor(0.0)) # The value 2, at any state.
	translated = 1.0 + speed * GRID # [5] the integral form x_t = x_0 + 2t from x_0 = 1.
	expanded = exact_linear(START_POINTS, torch.tensor(1.0)) # [3] states at t = 1, each doubled.
	contracted = exact_linear(START_POINTS, torch.tensor(1.0), a=-LOG2) # [3] states each halved.
	turns = torch.tensor([math.pi / 2, math.pi]) # [2] the quarter turn and the half turn.
	circle = torch.stack([torch.cos(turns), torch.sin(turns)], dim=-1) # [2, 2] x_t = (cos t, sin t).
	# The squared norm is conserved because the field is orthogonal to the position: 2 x . v_t(x) = 0.
	off_axis = torch.tensor([0.3, 0.7]) # [2] a state away from the axes, to test the conservation.
	norm_rate = 2.0 * (off_axis * rotation_field(off_axis, torch.tensor(0.0))).sum() # d\|\|x\|\|^2/dt.
	return {
	"constant_field_trajectory": [round(float(v), 4) for v in translated],
	"expansion_x_half_from_two": float(exact_linear(torch.tensor(2.0), torch.tensor(0.5))),
	"expansion_endpoints": [round(float(v), 4) + 0.0 for v in expanded],
	"expansion_outer_spacing": float(expanded[2] - expanded[0]),
	"contraction_endpoints": [round(float(v), 4) + 0.0 for v in contracted],
	"contraction_outer_spacing": float(contracted[2] - contracted[0]),
	"rotation_quarter_turn": [round(float(v), 4) + 0.0 for v in circle[0]],
	"rotation_half_turn": [round(float(v), 4) + 0.0 for v in circle[1]],
	"rotation_radius": float(circle.norm(dim=-1).mean()),
	"rotation_squared_norm_rate": float(norm_rate),
	}


	# %% 03. The flow map
	#
	# Section 2.4. We solve the ODE from every starting point at once. That gives the flow map phi_t.
	# At t = 0 that map is the identity, and for the expansion field it multiplies the whole line by
	# 2^t.
	def flow_map(x_0, t, a=LOG2):
	"""phi_t(x_0) = x_0 e^{at}, the accumulated transformation after following v_t for a time t."""
	return exact_linear(x_0, t, a) # [N] the state each starting point has reached by time t.


	def flow_map_report():
	"""The identity at t = 0, the Jacobian determinant at t = 1, and noncrossing trajectories."""
	probe = torch.linspace(-3.0, 3.0, 101) # [101] states spread over the line.
	identity_error = (flow_map(probe, torch.tensor(0.0)) - probe).abs().max() # phi_0 = id.
	x_0 = torch.tensor(1.3, requires_grad=True) # A state at which to differentiate phi_1.
	jacobian = torch.autograd.grad(flow_map(x_0, torch.tensor(1.0)), x_0)[0] # d phi_1 / d x_0 = 2.
	# Proposition 2.1: distinct trajectories of a Lipschitz field never share a state at a shared time.
	gaps = (flow_map(START_POINTS[2], GRID) - flow_map(START_POINTS[0], GRID)).abs() # [5].
	return {
	"flow_map_identity_error": float(identity_error),
	"flow_map_jacobian_det": float(jacobian),
	"trajectory_gap_minimum": float(gaps.min()),
	}


	# %% 04. Euler's method
	#
	# Section 2.5. Within one step the velocity stays fixed at its value at the start of the step. On
	# dx_t/dt = x_t the velocity grows with the state, so every step undershoots.
	def euler(field, x_0, steps):
	"""hat x_{k+1} = hat x_k + Delta t v_{t_k}(hat x_k) on a uniform grid of ``steps`` intervals."""
	x_hat = x_0.clone() # The running approximation, started from the exact initial state.
	dt = 1.0 / steps # Delta t, the width of one interval of the grid.
	for k in range(steps): # One velocity evaluation per step.
	x_hat = x_hat + dt * field(x_hat, torch.tensor(k * dt)) # The Euler update at grid time t_k.
	return x_hat


	def euler_ladder(largest):
	"""Euler on dx_t/dt = x_t from x_0 = 1, at every power of two up to ``largest`` steps."""
	exact = math.e # The exact solution x_1 = e^1, the chapter's 2.7183.
	counts, landings, errors = [], [], []
	steps = 1
	largest = max(largest, 4) # The chapter's three rungs, at one, two and four steps, always run.
	while steps <= largest: # 1, 2, 4, 8, ... up to the requested largest step count.
	landing = float(euler(lambda x, t: x, torch.tensor(1.0), steps)) # The field is v_t(x) = x.
	counts.append(steps)
	landings.append(round(landing, 6))
	errors.append(round(exact - landing, 6))
	steps *= 2
	return {
	"euler_step_counts": counts,
	"euler_landings": landings,
	"euler_errors": errors,
	# A first-order method has error proportional to 1/N, so the error times N settles at e/2.
	"euler_error_times_steps": [round(e * n, 4) for e, n in zip(errors, counts)],
	"euler_error_halving_ratios": [round(errors[i] / errors[i + 1], 4)
	for i in range(len(errors) - 1)],
	"euler_relative_error_four_steps": float((exact - landings[2]) / exact),
	}


	# %% 05. A random starting point
	#
	# Section 2.6. The initial state becomes a random variable. Each draw follows the same field, so
	# the law at time t is the pushforward p_t = (phi_t)_# p_0.
	def pushforward_moments(num_samples, seed):
	"""Moments of phi_1(X_0) for X_0 ~ N(0, 1) under the expansion field, against N(0, 4)."""
	set_seed(seed) # Same seed, same draw, same report.
	x_0 = torch.randn(num_samples) # [N] draws from the prior p_0 = N(0, 1).
	x_1 = flow_map(x_0, torch.tensor(1.0)) # [N] states at t = 1, every one of them doubled.
	return {
	"pushforward_mean": float(x_1.mean()), # Exact value 0, since the flow fixes the origin.
	"pushforward_variance": float(x_1.var(unbiased=False)), # Scaling by 2 gives variance 4.
	"pushforward_exact_variance": 4.0, # Scaling a centered law by c multiplies variance by c^2.
	}


	# %% 06. Probability flux and the continuity equation
	#
	# Section 2.7. The flux j_t(x) = p_t(x) v_t(x) measures probability transported per unit time.
	# Once an interval shrinks to a point, the balance over it becomes the continuity equation.
	def spreading_density(x, t):
	"""p_t(x) for X_t ~ N(0, 4^t), the law the expansion field produces from a standard normal."""
	std = torch.pow(torch.tensor(2.0), t) # 2^t, the standard deviation at time t.
	return torch.exp(-x ** 2 / (2.0 * std ** 2)) / (std * math.sqrt(2.0 * math.pi))


	def flux_balance():
	"""The two endpoint fluxes on [-1, 1] at t = 0, and the rate of change of the enclosed mass."""
	t_0 = torch.tensor(0.0)
	endpoints = torch.tensor([-1.0, 1.0]) # [2] the interval's endpoints a and b.
	density = spreading_density(endpoints, t_0) # [2] p_0(-1) and p_0(1), both 0.2420.
	flux = density * linear_field(endpoints, t_0) # [2] j_0(x) = p_0(x) v_0(x).
	# The same rate read off the distribution: M_t = erf(2^{-t} / sqrt 2), differentiated in t.
	t = torch.tensor(0.0, requires_grad=True)
	mass = torch.erf(torch.pow(torch.tensor(2.0), -t) / math.sqrt(2.0)) # P(\|X_t\| <= 1).
	mass_rate = torch.autograd.grad(mass, t)[0] # d M_t / dt at t = 0.
	return {
	"density_at_one": float(density[1]),
	"flux_at_minus_one": float(flux[0]),
	"flux_at_one": float(flux[1]),
	"flux_balance": float(flux[0] - flux[1]), # j_0(a) - j_0(b), the net rate into [-1, 1].
	"enclosed_mass": float(mass.detach()),
	"enclosed_mass_rate": float(mass_rate),
	}


	def continuity_check():
	"""Both sides of the continuity equation at x = 2 and t = 0, in the expanded product-rule form."""
	x = torch.tensor(2.0, requires_grad=True) # The position held fixed while the density changes.
	t = torch.tensor(0.0, requires_grad=True) # The time at which both sides are evaluated.
	p = spreading_density(x, t) # p_0(2), the chapter's 0.05399.
	dp_dt = torch.autograd.grad(p, t, create_graph=True)[0] # Local time derivative of the density.
	dp_dx = torch.autograd.grad(p, x, create_graph=True)[0] # Spatial gradient of the density.
	divergence = torch.autograd.grad(linear_field(x, t), x)[0] # div v = log 2 in one dimension.
	transport = linear_field(x, t) * dp_dx # v_t(x) . grad p_t(x), the transport term.
	volume = p * divergence # p_t(x) div v_t(x), the local volume-change term.
	return {
	"density_at_two": float(p.detach()),
	"density_gradient_at_two": float(dp_dx.detach()),
	"velocity_divergence": float(divergence),
	"continuity_time_derivative": float(dp_dt.detach()),
	"continuity_minus_flux_divergence": float(-(transport + volume).detach()),
	}


	# %% 07. The density along a moving trajectory
	#
	# Section 2.8. An observer riding with the sample sees no transport term, so the log density
	# changes at the rate -div v_t(x_t). Over a unit of time that rate agrees with the Jacobian of the
	# flow map.
	def log_density_along_trajectory(x_start=0.7):
	"""The instantaneous change-of-variables equation on the expansion field, both ways."""
	x_0 = torch.tensor(x_start) # One starting point, carried from t = 0 to t = 1.
	x_1 = flow_map(x_0, torch.tensor(1.0)) # The state the trajectory reaches at t = 1.
	log_p0 = torch.log(spreading_density(x_0, torch.tensor(0.0))) # log p_0(x_0).
	log_p1 = torch.log(spreading_density(x_1, torch.tensor(1.0))) # log p_1(x_1).
	origin = torch.tensor(0.0) # The fixed point of the flow, where phi_1(0) = 0.
	prior_at_origin = float(spreading_density(origin, torch.tensor(0.0))) # p_0(0) = 0.3989.
	final_at_origin = float(spreading_density(origin, torch.tensor(1.0))) # p_1(0) = 0.1995.
	return {
	"log_density_rate": -LOG2, # d log p_t(x_t) / dt = -div v_t(x_t), a constant here.
	"log_density_change": float(log_p1 - log_p0), # Measured from the two densities.
	"log_density_accumulated": -LOG2 * 1.0, # The integral of -div v_s over s from 0 to 1.
	"prior_density_at_origin": prior_at_origin,
	"final_density_at_origin": final_at_origin,
	"density_ratio_at_origin": prior_at_origin / final_at_origin, # \|det J_{phi_1}\| = 2.
	}


	# %% 08. The report
	def main():
	parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
	parser.add_argument("--steps", type=int, default=EULER_STEPS, help="largest Euler step count")
	parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="draws from the prior")
	parser.add_argument("--seed", type=int, default=0, help="random seed")
	parser.add_argument("--out", type=str, default="", help="directory for the run's report")
	parser.add_argument("--quiet", action="store_true", help="print only the final report")
	args = parser.parse_args()

	report = {}
	report.update(trajectories()) # Sections 2.2 and 2.3.
	report.update(flow_map_report()) # Section 2.4.
	report.update(euler_ladder(args.steps)) # Section 2.5.
	report.update(pushforward_moments(args.samples, args.seed)) # Section 2.6.
	report.update(flux_balance()) # Section 2.7.
	report.update(continuity_check()) # Section 2.7.
	report.update(log_density_along_trajectory()) # Section 2.8.

	if not args.quiet: # The Euler ladder, rung by rung, as the chapter develops it.
	for n, landing, error in zip(report["euler_step_counts"], report["euler_landings"],
	report["euler_errors"]):
	print(f" euler steps {n:3d} landing {landing:.6f} error {error:.6f}"
	f" error x steps {error * n:.4f}")

	print_report("velocity fields, flow maps and the continuity equation", report)
	save_run(args.out, report=report) # Writes report.json when --out names a directory.


	# %% 09. Run the script
	if __name__ == "__main__":
	main()