Download lecture_7/displacement.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 12.3 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/displacement.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_7/displacement.py
-
curl -L -o displacement.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/displacement.py
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() | |