CIS6270 / lecture_7 /displacement.py
pranamanam's picture
Upload 87 files
0ba9d09 verified
Raw History Blame Contribute Delete
12.3 kB
"""Transport as motion in time: the kinetic action, straight lines, and displacement interpolation.
A plan records which source point ends where and nothing about the trajectory between them. The
Benamou-Brenier formulation replaces the plan by a density and a velocity field and measures the
transport by the energy the particles expend. Its minimizer moves every particle in a straight line
at constant speed, with the endpoints paired by the optimal plan.
Chapter 7, Sections 7.7 to 7.7.5 of the notes. Runs on the one-particle example of Section 7.7.1
and on the Gaussian pair N(0,1) to N(3,4), both in closed form. Prints the action of a steady and a
dawdling particle, the two sides of the averaged bound, the interpolated density against the mixture
of the endpoints, and the velocity field with its action density.
python lecture_7/displacement.py # 1024 times on the grid
python lecture_7/displacement.py --steps 20 # a coarse time grid, for a first look
"""
# %% 01. Imports and configuration
from __future__ import annotations
import math
import torch
from cis6270.runner import common_parser, print_report, save_run, set_seed
DTYPE = torch.float64 # Double precision, so the closed forms match the chapter digit for digit.
MEAN_0, STD_0 = 0.0, 1.0 # The source N(0, 1).
MEAN_1, STD_1 = 3.0, 2.0 # The target N(3, 4), with standard deviation 2.
# %% 02. One particle
#
# The squared displacement of a particle is at most the integral of its squared speed, by
# Cauchy-Schwarz against the constant function 1 on an interval of length 1, with equality when the
# velocity is constant in time. Section 7.7.1.
def action_of_a_path(positions: torch.Tensor, times: torch.Tensor) -> float:
"""The integral of the squared speed along a trajectory sampled at the given times."""
velocity = (positions[1:] - positions[:-1]) / (times[1:] - times[:-1]) # [n-1] piecewise speed.
return float((velocity ** 2 * (times[1:] - times[:-1])).sum()) # Integral of ||x_dot||^2.
def steady_particle(times: torch.Tensor) -> torch.Tensor:
"""A particle covering the distance 2 at the constant speed 2, so its action is 4."""
return 2.0 * times # x(t) = 2t.
def dawdling_particle(times: torch.Tensor) -> torch.Tensor:
"""A particle that waits until t = 1/2 and then covers the same distance at speed 4."""
return torch.clamp(4.0 * (times - 0.5), min=0.0) # Zero on [0, 1/2], then 4(t - 1/2).
# %% 03. Averaging the bound over the particles
#
# The joint law of the two endpoints of a particle is a coupling of mu and nu, so the action is at
# least the average squared displacement under that coupling, and that is at least W_2^2.
# Section 7.7.2.
def average_squared_displacement(x: torch.Tensor, y: torch.Tensor, pi: torch.Tensor) -> float:
"""The average of ||x_1 - x_0||^2 over endpoint pairs drawn from the coupling pi."""
return float((pi * (x[:, None] - y[None, :]) ** 2).sum()) # Sum over routes of mass times cost.
# %% 04. Straight lines
#
# We draw the endpoint pair from the optimal coupling and move along the segment between them.
# The velocity is constant along it. Cauchy-Schwarz is then tight for every particle, and the
# coupling is optimal, so the second inequality is tight as well. Section 7.7.3.
def straight_path(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""x(t) = (1 - t) x_0 + t x_1, with the constant velocity x_1 - x_0."""
return (1.0 - t) * x0 + t * x1 # Same shape as x0 and x1 broadcast against t.
# %% 05. Displacement interpolation
#
# We read off the law of the particles at an intermediate time and get the displacement
# interpolation. Under the optimal map X_1 = 3 + 2 X_0 of the Gaussian example, X_t collects into
# 3t + (1 + t) X_0, so rho_t is Gaussian with mean 3t and standard deviation 1 + t. Section 7.7.4.
def displacement_law(t: float) -> tuple[float, float]:
"""The mean and standard deviation of rho_t for the Gaussian pair, both linear in t."""
return 3.0 * t, 1.0 + t # Mean slides from 0 to 3, standard deviation from 1 to 2.
def normal_density(x: torch.Tensor, mean: float, std: float) -> torch.Tensor:
"""The density of N(mean, std^2) evaluated at every entry of x."""
return torch.exp(-0.5 * ((x - mean) / std) ** 2) / (std * math.sqrt(2.0 * math.pi))
def mixture_density(x: torch.Tensor) -> torch.Tensor:
"""Half the source density plus half the target density, the other path between the two ends."""
return 0.5 * normal_density(x, MEAN_0, STD_0) + 0.5 * normal_density(x, MEAN_1, STD_1)
def gaussian_w2(m0: float, s0: float, m1: float, s1: float) -> float:
"""W_2 between two Gaussians on the line, the closed form of Section 7.6."""
return math.sqrt((m0 - m1) ** 2 + (s0 - s1) ** 2)
# %% 06. The velocity field behind the displacement path
#
# Each particle obeys X_t = 3t + (1 + t) X_0 with the constant velocity 3 + X_0. We solve the
# trajectory for X_0 and substitute. The velocity is then a function of the current position
# alone. Section 7.7.5.
def velocity_field(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""v_t(x) = 3 + (x - 3t) / (1 + t), the mean speed plus a stretch about the mean."""
return 3.0 + (x - 3.0 * t) / (1.0 + t) # Broadcasts over a grid of states and times.
def continuity_residual(x: float, t: float) -> float:
"""The left side of d_t rho_t + d_x (rho_t v_t) = 0, computed by automatic differentiation."""
xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
ts = torch.tensor(t, dtype=DTYPE, requires_grad=True)
rho = normal_density(xs, 3.0 * ts, 1.0 + ts) # rho_t written with the two linear schedules.
flux = rho * velocity_field(xs, ts) # The probability flux rho_t v_t.
d_rho_dt = torch.autograd.grad(rho, ts, retain_graph=True)[0] # Keep the graph for the flux.
d_flux_dx = torch.autograd.grad(flux, xs)[0]
return float((d_rho_dt + d_flux_dx).detach())
def action_density(t: float, levels: int) -> float:
"""The integral of rho_t(x) v_t(x)^2 over x, or E[(3 + X_0)^2] at every time."""
mean, std = displacement_law(t)
q = (torch.arange(levels, dtype=DTYPE) + 0.5) / levels # Quantile levels of rho_t.
x = mean + std * torch.special.ndtri(q) # [levels] a stratified sample of rho_t.
speed = velocity_field(x, torch.tensor(t, dtype=DTYPE)) # [levels].
return float((speed ** 2).mean()) # The average of the squared speed under rho_t.
# %% 07. The report
def main() -> None:
parser = common_parser(__doc__.split("\n")[0])
parser.set_defaults(steps=1024) # --steps is the number of times on the grid.
args = parser.parse_args()
set_seed(args.seed)
loud = not args.quiet
n = max(args.steps, 2)
times = torch.linspace(0.0, 1.0, n + 1, dtype=DTYPE) # [n+1] the time grid.
steady = action_of_a_path(steady_particle(times), times)
dawdle = action_of_a_path(dawdling_particle(times), times)
if loud:
print(f"One particle from 0 to 2 on a grid of {n} steps")
print(f" constant speed 2: action {steady:.3f}, squared displacement {2.0 ** 2:.0f}")
print(f" waiting then speed 4: action {dawdle:.3f}, the same two endpoints")
x = torch.tensor([0.0, 2.0], dtype=DTYPE) # Four points, mass 1/2 each.
y = torch.tensor([1.0, 3.0], dtype=DTYPE)
ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=DTYPE)
crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=DTYPE)
independent = torch.full((2, 2), 0.25, dtype=DTYPE)
w22_points = average_squared_displacement(x, y, ordered) # The optimal coupling on the line.
if loud:
print(f"\nAveraging over particles, four points with mass 1/2")
for name, pi in (("ordered", ordered), ("independent", independent), ("crossed", crossed)):
print(f" {name:<12s} E||x_1 - x_0||^2 = "
f"{average_squared_displacement(x, y, pi):.1f} >= W_2^2 = {w22_points:.1f}")
# Straight lines under the optimal coupling, with the endpoints of the Gaussian example.
q = (torch.arange(n, dtype=DTYPE) + 0.5) / n
x0 = MEAN_0 + STD_0 * torch.special.ndtri(q) # [n] a stratified sample of N(0, 1).
x1 = MEAN_1 + (STD_1 / STD_0) * (x0 - MEAN_0) # The Brenier map 3 + 2x, the optimal pairing.
paths = straight_path(x0[None, :], x1[None, :], times[:, None]) # [n+1, n] every trajectory.
speeds = (paths[1:] - paths[:-1]) / (times[1:] - times[:-1])[:, None] # [n, n] piecewise speed.
action = float((speeds ** 2 * (times[1:] - times[:-1])[:, None]).sum(dim=0).mean())
squared_displacement = float(((x1 - x0) ** 2).mean())
if loud:
print("\nStraight lines under the optimal coupling")
print(f" action averaged over the particles {action:.4f} "
f"average squared displacement {squared_displacement:.4f}")
mean_half, std_half = displacement_law(0.5)
grid = torch.linspace(-6.0, 9.0, 15001, dtype=DTYPE) # A dense grid for the two densities.
mixture = mixture_density(grid)
peak = float(grid[int(mixture.argmax())])
interpolated = normal_density(grid, mean_half, std_half)
separation = float(0.5 * (interpolated - mixture).abs().sum() * (grid[1] - grid[0]))
half_distance = gaussian_w2(MEAN_0, STD_0, mean_half, std_half)
whole_distance = gaussian_w2(MEAN_0, STD_0, MEAN_1, STD_1)
if loud:
print(f"\nAt t = 1/2 the displacement interpolation is N({mean_half:.1f}, "
f"{std_half ** 2:.2f}) with standard deviation {std_half:.1f}")
print(f" the mixture of the two ends peaks at x = {peak:.2f} and is not a Gaussian")
print(f" total variation between the two candidate densities {separation:.3f}")
print(f" W_2(rho_0, rho_1/2) = {half_distance:.3f} = "
f"{half_distance / whole_distance:.1f} x W_2(mu, nu) = {whole_distance:.3f}")
mean_speed = float(velocity_field(torch.tensor(3.0 * 0.5, dtype=DTYPE),
torch.tensor(0.5, dtype=DTYPE))) # A particle at the mean.
above_speed = float(velocity_field(torch.tensor(3.0 * 0.5 + 1.5, dtype=DTYPE),
torch.tensor(0.5, dtype=DTYPE))) # One standard deviation up.
residual = max(abs(continuity_residual(xx, tt))
for xx in (-1.0, 0.0, 1.5, 4.0) for tt in (0.1, 0.5, 0.9))
densities = [action_density(float(t), 2048) for t in torch.linspace(0.0, 1.0, 5, dtype=DTYPE)]
landing = x0 + 1.0 * velocity_field(x0, torch.zeros((), dtype=DTYPE)) # One Euler step of size 1.
one_step = float((landing - x1).abs().max()) # The step lands on the Brenier map 3 + 2 x_0.
if loud:
print(f"\nv_t(x) = 3 + (x - 3t) / (1 + t)")
print(f" at the mean, speed {mean_speed:.0f}; one standard deviation above, "
f"speed {above_speed:.0f}")
print(f" action density at five times {[round(d, 4) for d in densities]}")
print(f" largest continuity residual on the test states {residual:.2e}")
report = {
"action of the steady particle": steady,
"action of the dawdling particle": dawdle,
"squared displacement of both": 4.0,
"E||x1-x0||^2, ordered": average_squared_displacement(x, y, ordered),
"E||x1-x0||^2, independent": average_squared_displacement(x, y, independent),
"E||x1-x0||^2, crossed": average_squared_displacement(x, y, crossed),
"action of the straight paths": action,
"W_2^2 of the Gaussian pair": whole_distance ** 2,
"rho_1/2 mean, standard deviation": f"{mean_half:.1f}, {std_half:.1f}",
"peak of the mixture": peak,
"total variation of the two paths": separation,
"W_2(rho_0, rho_1/2)": half_distance,
"W_2(mu, nu)": whole_distance,
"speed at the mean": mean_speed,
"speed one deviation above": above_speed,
"action density, every time": max(densities),
"continuity equation residual": residual,
"one Euler step of size 1 minus T": one_step,
}
print_report("Transport as motion, Chapter 7 Sections 7.7 to 7.7.5", report)
save_run(args.out, config=vars(args), report=report) # Written only with --out.
if __name__ == "__main__":
main()