Delete lecture_7

lecture_7/README.md

@@ -1,166 +0,0 @@
-# Lecture 7: Optimal Transport and Schrodinger Bridges
-
-Chapter 7 of the notes.
-
-Every model of Chapters 2 through 6 trains on a pair of samples, and in almost all of them the two
-were drawn independently. This lecture takes the pairing as the quantity to be chosen. The first
-four files price a pairing, relax Monge's map to Kantorovich's coupling, certify the optimum with
-prices, smooth the problem with an entropy term that Sinkhorn's algorithm solves by rescaling rows
-and columns, and read transport as motion in time, where the minimizer moves every particle in a
-straight line. The next five replace the cost matrix by the endpoint law of a reference process,
-which turns the same scaling into a Schrodinger bridge: the chain rule puts reference bridges
-between every pair of ends, Doob's h-transform makes the bridge a process we can run forward, and
-three learning methods differ only in what they simulate and what they regress. Three more carry
-the construction to a finite state space, where a rate matrix replaces the drift, where a molecular
-graph factorizes into forty-five categorical features, and where discrete time makes the whole path
-law a table of sixteen numbers. The last file covers a population that splits into weighted branches
-and a system of interacting particles whose own dynamics is the reference.
-
-Almost every number below is computed in closed form, which is what the chapter does. Two files
-train.
-
-| file | command | what it computes |
-| ---- | ------- | ---------------- |
-| `optimal_transport.py` | `python lecture_7/optimal_transport.py` | the four points, the two-by-two plan and its prices, the sorted pairing on three points, and the Gaussian map |
-| `sinkhorn.py` | `python lecture_7/sinkhorn.py` | the Gibbs kernel, the first sweep by hand, the converged entropic plan, and the plan at four values of epsilon |
-| `displacement.py` | `python lecture_7/displacement.py` | the kinetic action of one particle and of a population, the displacement interpolation, and its velocity field |
-| `ot_flow_matching.py` | `python lecture_7/ot_flow_matching.py` | **trains** two velocity networks on the four-Gaussian target, one per coupling, and measures the straightness of each |
-| `static_bridge.py` | `python lecture_7/static_bridge.py` | the eight path probabilities, the chain rule, proportional fitting, the tilted matrices, and the bridge path law |
-| `dynamic_bridge.py` | `python lecture_7/dynamic_bridge.py` | Girsanov on a constant control, the backward and Hamilton-Jacobi-Bellman residuals, and the Brownian bridge |
-| `dsb.py` | `python lecture_7/dsb.py` | the two half-steps of proportional fitting on path laws, and the diffusion reversal on a Gaussian |
-| `dsbm.py` | `python lecture_7/dsbm.py` | the averaged bridge drift at one point, and one round of Markovian fitting on the chain |
-| `sf2m.py` | `python lecture_7/sf2m.py` | the two closed-form regression targets at a bridge sample, and the entropic plan of a batch |
-| `discrete_bridge.py` | `python lecture_7/discrete_bridge.py` | the two-state continuous-time bridge, its Euler sampler and its reverse rates, and the cost computed twice |
-| `ddsbm.py` | `python lecture_7/ddsbm.py` | the single bond with three types, the rates pinned at a terminal graph, and the counting over forty-five features |
-| `csbm.py` | `python lecture_7/csbm.py` | the two projections in discrete time, and D-IMF on sixteen paths against the reference's mixing |
-| `branched.py` | `python lecture_7/branched.py` | the two-branch energies, the tilt of a reference, and **trains** a bias force on two interacting particles |
-
-Every file takes `--seed` and `--quiet`, and `--steps` means the quantity that file iterates: the
-number of quantile levels, Sinkhorn sweeps, times on a grid, optimizer steps, fitting rounds or
-Euler steps. The slowest default run is about seven seconds; all thirteen run on a laptop CPU with
-no download.
-
-## What a default run produces
-
-`optimal_transport.py`. The ordered pairing of the four points costs 1 against 5 for the crossed
-one and 3 for the independent coupling. The two-by-two optimum is `[[0.3, 0.3], [0, 0.4]]` at cost
-3.4, the prices `f = (0, -8)` and `g = (1, 9)` have slack `[[0, 0], [8, 0]]` and dual value 3.4, so
-the duality gap is zero and `W_2 = 1.8439`. The six permutations of the three-point example cost
-11/3, 17/3, 35/3, 15, 77/3 and 27, and the quantile integral on 3072 levels returns 3.6667.
-Transporting N(0,1) to N(3,4) gives `W_2^2 = 10` under the map `T(x) = 3 + 2x`, and the same map on
-a stratified sample gives 9.9996.
-
-`sinkhorn.py`. The kernel is `[[0.606531, 0.011109], [0.606531, 0.606531]]`. The first sweep gives
-`u = (0.971440, 0.329744)` and `v = (0.38013, 3.32081)`, and the marginal error falls 0.680, 0.382,
-0.155, 0.047, 0.012, reaching 1e-9 after 17 sweeps. The entropic plan is
-`[[0.293122, 0.306878], [0.006878, 0.393122]]` at cost 3.45502 against 3.4, and the regularized
-objective is 1.208 at it, 1.222 at the exact optimum and 1.792 at independence. At epsilon = 10,
-2, 0.5 and 0.1 the cost is 4.0540165, 3.4550204, 3.4000004 and 3.4000000, in 7, 17, 28 and 80
-sweeps.
-
-`displacement.py`. A steady particle has action 4 and a dawdling one 8, for the same endpoints. The
-average squared displacement is 1, 3 and 5 under the ordered, independent and crossed pairings. At
-t = 1/2 the displacement interpolation is N(1.5, 2.25) and the mixture of the two ends peaks at
-0.13, with total variation 0.257 between them, and `W_2(rho_0, rho_1/2) = 1.5811`, half of 3.1623.
-The straight paths have action 9.9987, which is their average squared displacement. The action
-density is 9.9994 at every time on the grid, the stratified estimate of the exact 10, and the
-continuity residual is 1.1e-16.
-
-`ot_flow_matching.py`. On the two-point batch the independent coupling has average squared velocity
-3 and the optimal one 1. Over training the batch coupling cost is 10.31 against 3.25, the learned
-paths have straightness 1.80 against 0.065, and the one-step error is 2.56 against 0.41. Both
-models find all four modes. At `--steps 1500`, about eighteen seconds, the optimal coupling reaches
-straightness 0.024 and one-step error 0.28 while the independent one does not improve.
-
-`static_bridge.py`. The eight reference paths sum to one with `R(1,1,2) = 0.105`. The chain rule
-splits 0.2180 into 0.0201 and 0.1979. The reference lands 0.73 away from the requested terminal
-law, and proportional fitting reaches `f = (1.072002, 0.937019)` and `g = (0.350517, 1.843637)`,
-with the terminal error falling by 187.5 per round. The bridge coupling is
-`[[0.114605, 0.385395], [0.085395, 0.414605]]` at `KL = 0.2820` against 0.2847 for independence.
-The tilted matrices carry (0.5, 0.5) to (0.449213, 0.550787) to (0.2, 0.8), the path ratios take
-four values, and `KL(P* || R) = 0.281961`.
-
-`dynamic_bridge.py`. The relative entropy of a constant control u = 2 is 2 from the control energy,
-2 from the two endpoint laws and 1.9965 from the simulated paths. The backward residual is 2.2e-16
-and the Hamilton-Jacobi-Bellman residual 8.9e-16, and `-grad V` equals the bridge drift to 4.4e-16.
-At t = 0.25 the bridge has mean 0.5 and variance 0.1875, and at the state 0.6732 its drift is
-1.7691 against the straight-line velocity 2. At t = 1/2 the simulated mean is 1.0095 and the
-simulated variance 0.2563, against the exact 1 and 0.25.
-
-`dsb.py`. The reversal of the last step is `[[0.681416, 0.318584], [0.379310, 0.620690]]`. The
-backward chain started from (0.2, 0.8) reaches (0.466585, 0.533415), which is 0.0668 from the source
-law, and the second half-step returns the terminal law (0.201967, 0.798033), off by 0.0039. That is
-the round-1 error of proportional fitting on the coupling, and the error falls by 187.5 per round.
-Run backward from its own terminal law, the diffusion reversal lands at variance 0.9836 against the
-exact 1.
-
-`dsbm.py`. Two bridges through x = 1 at t = 0.5 have drifts 4 and -2 and average 1, where the
-derivative of the squared error vanishes. One round of Markovian fitting gives
-`M_0 = [[0.591425, 0.408575], [0.307692, 0.692308]]` and
-`M_1 = [[0.298457, 0.701543], [0.119588, 0.880412]]`, whose endpoint coupling has exact marginals
-and brings the gap from 0.0584 to 0.0077. The gap then falls by 7.37 per round and nine rounds reach
-1e-9.
-
-`sf2m.py`. At t = 0.25 with z = 0.4 the sample is 0.6732, the drift 1.7691, the score -0.9238, the
-flow target 2.2309 and the scaled score exactly -0.4. At t = 0.75 the coefficient changes sign and
-the flow target is 1.7691, so the two lie symmetrically about the straight-line velocity 2.
-
-`discrete_bridge.py`. Halving every rate costs 0.1534 and doubling costs 0.3863. The two-state
-bridge has `g = (0.3349, 1.9643)`, `phi_1/2 = (0.7999, 1.2667)` and the generator
-`[[-1.5836, 1.5836], [0.9472, -0.9472]]` at t = 1/2, where the marginal is (0.4668, 0.5332) and the
-reverse rates 1.3862 and 1.0821 match the forward flux at 0.7392. The Euler sampler's terminal
-error is 0.0223, 0.0106, 0.0051 and 0.0020 as the step is halved twice and then taken to a hundred
-steps. The cost is 0.3235 on the endpoints and 0.3235 on the paths, with the integrand 0.1300 at
-t = 1/2, and it rises 0.3148, 0.3235, 0.3337, 0.3348 toward the limit 0.3348 as the reference
-speeds up.
-
-`ddsbm.py`. The bond reference has `beta = 0.9163` and off-diagonal rate 0.3054, and the first
-column scaling is (0.2083, 0.7143, 2.9167). Two hundred sweeps give `f = (1.1779, 0.9506, 0.5062)`
-and `g = (0.1842, 0.6918, 3.0004)`, a plan that moves 0.4948 of the mass from no bond to double and
-leaves 0.2612 unchanged, between 0.18 for independence and 0.4 for the edit-minimizing plan. At
-r = 0.2, 0.4, 0.6, 0.8 and 0.99 the mass that stays is 0.2250, 0.2612, 0.2926, 0.3248 and 0.3812.
-The survival ratio reaches 1/2 at t = 0.2435, where the pinned rates are 1.2217, 0.0764 and 0.3054,
-a factor of 16 between a move onto the terminal type and a move off it, and one such edge against a
-network rate of 1 contributes 0.0229. With two terminal types the network learns 0.8552 at a loss
-of 0.1357, against 0.1467 at 1 and 0.2395 at 0.5. Nine heavy atoms give 45 features and 9.2e+27
-graphs, at a total jump rate of 24.74 per unit time.
-
-`csbm.py`. With alpha = 0.35 and three transitions, `A^3 = [[0.5135, 0.4865], [0.4865, 0.5135]]`
-and the exact bridge coupling is `[[0.10432, 0.39568], [0.09568, 0.40432]]`. The path 1,1,1,1 has
-reference bridge weight 0.5348, so its probability is 0.0535 at the start and 0.0558 under the
-bridge. The two laws agree at the interior times, (0.4754, 0.5246) and (0.4107, 0.5893), and differ
-by 0.0173 over the paths. One round of D-IMF brings that to 2.9e-04 and later rounds divide it by
-58.9, reaching 2.4e-11 after five. The last conditional (0.3160, 0.6840) against a network at
-(0.25, 0.75) costs 0.0110 nats. At alpha = 0.2, 0.35 and 0.4 the starting gap is 0.13589, 0.01728
-and 0.00512 and the rounds divide it by 5, 59 and 284; on 8, 16 and 64 paths it is 0.05743, 0.01728
-and 0.00156 and the rounds divide it by 7, 59 and 4674. The ordered reference out of the middle of
-five categories is (0.0545, 0.2442, 0.4026, 0.2442, 0.0545).
-
-`branched.py`. The two branches have weighted energies 1.7 and 0.3 for a total of 2.0, and with the
-second target at -4 the linear schedule costs 2.9 against 2.6 for one that releases mass late. The
-bent interpolant has energy 2.1667 against 2 for the straight line. The two-outcome tilt has
-`Z = 0.109`, raises the chance of reaching B from 0.1 to 0.917, and costs 1.836 nats, with
-`J(b*) = -log Z = 2.216`. Training the bias force raises the mean terminal reward from -21.9 under
-the reference to -6.9. A longer run does not improve that, because the importance weights
-concentrate as the control grows and the effective sample size falls with it.
-
-## What is simplified
-
-- Every example runs on the chapter's own two states, three bond types, four points or two
-  particles. Nothing here is a benchmark, and no file touches a real dataset.
-- The minibatch coupling of `ot_flow_matching.py` is an exact assignment through
-  `scipy.optimize.linear_sum_assignment`, which needs the two sides to be the same size. The
-  chapter's general `n` by `m` problem with unequal masses appears only in `sinkhorn.py`.
-- The two training files use small networks and short runs, so their reports show the direction of
-  an effect and not its converged size.
-- `dsb.py` and `dsbm.py` carry out their projections exactly on the full table of paths. The
-  networks the chapter describes, which replace those projections by a regression when the state
-  space is large, are not trained here.
-- `discrete_bridge.py` and `ddsbm.py` work one categorical feature at a time. `ddsbm.py` counts the
-  forty-five features of a molecular graph and the rate of the product process, and no graph is
-  built.
-- `csbm.py` runs both projections exactly on the sixteen-entry table. The network that replaces the
-  conditionals, and the simulation that replaces the exact endpoint coupling, are not trained here.
-- `branched.py` uses constant growth rates taken from the worked example, and trains only the bias
-  force of the entangled construction. The four-stage training of the branched construction and the
-  state cost `V_t` are not fitted.

lecture_7/branched.py

@@ -1,344 +0,0 @@
-"""Two bridges that change the shape of the problem: branching, and interacting particles.
-
-A population leaving one source does not always end in one place, and Branched Schrodinger Bridge
-Matching gives each outcome its own controlled diffusion and its own share of the mass, with a
-learned growth rate moving mass between branches. Entangled Schrodinger Bridge Matching keeps one
-population and changes the reference, taking the molecular dynamics of the system itself and
-tilting it at the final time by a target density, which a bias force on every particle reproduces.
-
-Chapter 7, Sections 7.26 to 7.27.3 of the notes. Runs the two-branch example in closed form and
-trains a bias force on two interacting particles in the plane. Prints the weighted energies of the
-two branch schedules, the interpolant and its energy, the reweighting arithmetic of the two-outcome
-example, one Euler-Maruyama step, and the reward the trained force reaches.
-
-    python lecture_7/branched.py              # 100 optimizer steps on the bias force
-    python lecture_7/branched.py --steps 20   # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import math
-
-import torch
-import torch.nn as nn
-
-from cis6270.nets import TimeEmbedding, count_parameters
-from cis6270.runner import (
-    common_parser, pick_device, print_report, save_run, set_seed, train,
-)
-
-PARTICLES = 2  # n, the number of interacting particles.
-DIMENSION = 2  # Each particle lives in the plane, so the projector below has something to remove.
-MASS, FRICTION, TEMPERATURE = 1.0, 1.0, 1.0  # m_i, gamma and k_B tau, all in reduced units.
-NOISE = math.sqrt(2.0 * FRICTION * TEMPERATURE / MASS)  # The noise scale of the reference.
-TARGET_WIDTH = 0.4  # The width of the target density around the region B.
-PATHS = 64  # Trajectories simulated per optimizer step.
-SIM_STEPS = 30  # Euler-Maruyama steps per trajectory.
-
-
-# %% 02. A population that splits
-#
-# Branch k is a controlled diffusion carrying the fraction w^k_t of the population, and the weights
-# change at a learned growth rate with sum_k g^k_t = 0, so a branch grows only by taking mass from
-# another. Section 7.26.
-
-
-def branch_weights(growth, t: torch.Tensor) -> torch.Tensor:
-    """w^1_t = integral of the growth rate from 0 to t, with branch 0 holding the remainder."""
-    return growth(t)  # The weight of branch 1; branch 0 carries 1 - w^1_t.
-
-
-def weighted_energy(drift: float, weight, points: int = 2001) -> float:
-    """The time integral of (1/2) ||u||^2 times the weight of the branch, by the trapezoid rule."""
-    t = torch.linspace(0.0, 1.0, points)
-    density = 0.5 * drift ** 2 * weight(t)  # The control energy priced by the mass that spends it.
-    return float(torch.trapezoid(density, t))
-
-
-def branch_sample(weight_one: float, generator: torch.Generator) -> int:
-    """Draw a branch with probability w^k_t, which is how a sample of the population is taken."""
-    return int(torch.rand(1, generator=generator) < weight_one)  # 1 for branch 1, 0 for branch 0.
-
-
-# %% 03. The interpolant
-#
-# A network correction multiplied by t(1 - t) leaves both endpoints fixed for every value of the
-# parameters, so the endpoint constraints hold during training with no penalty term. With no state
-# cost the energy of a bent path exceeds the energy of the straight line, and the training drives
-# the correction to zero. Section 7.26.2.
-
-
-def interpolant(x0: float, x1: float, correction: float, t: torch.Tensor) -> torch.Tensor:
-    """x_t = (1 - t) x_0 + t x_1 + t(1 - t) phi, with phi here a constant."""
-    return (1.0 - t) * x0 + t * x1 + t * (1.0 - t) * correction
-
-
-def interpolant_velocity(x0: float, x1: float, correction: float, t: torch.Tensor) -> torch.Tensor:
-    """x_dot = (x_1 - x_0) + (1 - 2t) phi, the derivative of the interpolant in t."""
-    return (x1 - x0) + (1.0 - 2.0 * t) * correction
-
-
-def interpolant_energy(x0: float, x1: float, correction: float, points: int = 2001) -> float:
-    """The integral of one half the squared velocity of the interpolant over the unit interval."""
-    t = torch.linspace(0.0, 1.0, points)
-    return float(torch.trapezoid(0.5 * interpolant_velocity(x0, x1, correction, t) ** 2, t))
-
-
-# %% 04. The reference of Entangled Schrodinger Bridge Matching
-#
-# The reference is the dynamics the system actually follows, underdamped Langevin with the physical
-# force, the friction and the thermal noise. Section 7.27.
-
-
-def potential(positions: torch.Tensor) -> torch.Tensor:
-    """U(R), a double well on each particle plus a spring that couples the two. [B, n, d] -> [B]."""
-    radial = ((positions ** 2).sum(dim=-1) - 1.0) ** 2  # [B, n] a well at every unit circle.
-    coupling = 0.5 * ((positions[:, 0] - positions[:, 1]) ** 2).sum(dim=-1)  # [B] the spring.
-    return radial.sum(dim=-1) + coupling  # [B].
-
-
-def physical_force(positions: torch.Tensor) -> torch.Tensor:
-    """-grad_R U(R), differentiated by hand so the simulation needs no autograd. [B, n, d]."""
-    radial = -4.0 * ((positions ** 2).sum(dim=-1, keepdim=True) - 1.0) * positions  # [B, n, d].
-    separation = positions[:, 0] - positions[:, 1]  # [B, d] the displacement across the spring.
-    spring = torch.stack([-separation, separation], dim=1)  # [B, n, d] equal and opposite.
-    return radial + spring
-
-
-def force_check(positions: torch.Tensor) -> float:
-    """The largest gap between the force written by hand and the gradient autograd computes."""
-    r = positions.detach().requires_grad_(True)
-    gradient = torch.autograd.grad(potential(r).sum(), r)[0]
-    return float((physical_force(positions) + gradient).abs().max())
-
-
-TARGET = torch.tensor([1.0, 0.0])  # The centre of the region B, the same for both particles.
-
-
-def log_target(positions: torch.Tensor) -> torch.Tensor:
-    """r(X_1) = log pi_B(R_1), zero inside the region and negative outside. [B, n, d] -> [B]."""
-    offset = positions - TARGET.to(positions)  # [B, n, d] displacement from the target.
-    return -(offset ** 2).sum(dim=(1, 2)) / (2.0 * TARGET_WIDTH ** 2)  # [B].
-
-
-def target_direction(positions: torch.Tensor) -> torch.Tensor:
-    """s_hat_i, the unit vector along the gradient of log pi_B for particle i. [B, n, d]."""
-    gradient = -(positions - TARGET.to(positions)) / TARGET_WIDTH ** 2  # [B, n, d].
-    return gradient / gradient.norm(dim=-1, keepdim=True).clamp(min=1e-8)
-
-
-# %% 05. The bias force and its parameterization
-#
-# A nonnegative scalar along the direction of steepest increase of the target density, plus a second
-# output with its component along that direction removed, so the push toward the target is never
-# negative and is read off directly as alpha. Section 7.27.3.
-
-
-def compose_bias(alpha: torch.Tensor, h: torch.Tensor, s_hat: torch.Tensor) -> torch.Tensor:
-    """b = alpha s_hat + (I - s_hat s_hat^T) h, the parameterization of the force on each particle."""
-    along = (h * s_hat).sum(dim=-1, keepdim=True) * s_hat  # [B, n, d] the removed component.
-    return alpha.unsqueeze(-1) * s_hat + (h - along)  # [B, n, d].
-
-
-class BiasForce(nn.Module):
-    """A network from the positions, the velocities and the time to one force per particle."""
-
-    def __init__(self, width: int = 64):
-        super().__init__()
-        self.time_embedding = TimeEmbedding(width, scale=1.0)
-        self.input = nn.Linear(2 * PARTICLES * DIMENSION, width)  # Positions and velocities.
-        self.body = nn.Sequential(nn.Linear(width, width), nn.SiLU(),
-                                  nn.Linear(width, width), nn.SiLU())
-        self.head = nn.Linear(width, PARTICLES * (1 + DIMENSION))  # alpha and h per particle.
-
-    def forward(self, positions: torch.Tensor, velocities: torch.Tensor,
-                t: torch.Tensor) -> torch.Tensor:
-        """positions, velocities: [B, n, d], t: [B] -> the bias force [B, n, d]."""
-        state = torch.cat([positions.flatten(1), velocities.flatten(1)], dim=1)  # [B, 2nd].
-        hidden = self.body(self.input(state) + self.time_embedding(t))  # [B, width].
-        output = self.head(hidden).view(-1, PARTICLES, 1 + DIMENSION)  # [B, n, 1 + d].
-        alpha = nn.functional.softplus(output[..., 0])  # [B, n] nonnegative, so b never points away.
-        return compose_bias(alpha, output[..., 1:], target_direction(positions))
-
-
-# %% 06. Simulating the controlled dynamics
-#
-# The control enters the velocity equation divided by the mass, against the noise scale, so in units
-# of the noise it is b / (m sigma) and Girsanov prices it at one half its squared size.
-# Sections 7.27.1 and 7.27.2.
-
-
-@torch.no_grad()
-def roll_out(force: BiasForce | None, paths: int, steps: int, device: torch.device,
-             generator: torch.Generator) -> dict:
-    """Run the Langevin dynamics with the given bias force, storing the states and the increments."""
-    dt = 1.0 / steps
-    positions = torch.full((paths, PARTICLES, DIMENSION), -1.0, device=device)
-    positions[:, :, 1] = 0.0  # Both particles start at the far well.
-    velocities = torch.randn(paths, PARTICLES, DIMENSION, generator=generator).to(device) * \
-        math.sqrt(TEMPERATURE / MASS)  # Maxwell-Boltzmann velocities at the reference temperature.
-    stored_positions, stored_velocities, stored_u, increments = [], [], [], []
-    energy = torch.zeros(paths, device=device)  # (1/2) int ||u||^2 dt along each path.
-    cross = torch.zeros(paths, device=device)  # int u . dW along each path.
-    for step in range(steps):
-        t = torch.full((paths,), step * dt, device=device)
-        u = (force(positions, velocities, t) / (MASS * NOISE)) if force is not None else \
-            torch.zeros_like(positions)  # The control in units of the noise.
-        noise = torch.randn(paths, PARTICLES, DIMENSION,
-                            generator=generator).to(device) * math.sqrt(dt)
-        stored_positions.append(positions.clone())
-        stored_velocities.append(velocities.clone())
-        stored_u.append(u.clone())
-        increments.append(noise.clone())
-        energy = energy + 0.5 * (u ** 2).sum(dim=(1, 2)) * dt
-        cross = cross + (u * noise).sum(dim=(1, 2))
-        acceleration = physical_force(positions) / MASS - FRICTION * velocities  # The reference.
-        velocities = velocities + (acceleration + NOISE * u) * dt + NOISE * noise
-        positions = positions + velocities * dt  # The position follows the velocity.
-    return {
-        "positions": torch.stack(stored_positions), "velocities": torch.stack(stored_velocities),
-        "u_bar": torch.stack(stored_u), "noise": torch.stack(increments),
-        "reward": log_target(positions), "energy": energy, "cross": cross, "dt": dt,
-    }
-
-
-def importance_weights(rollout: dict) -> torch.Tensor:
-    """softmax over the batch of r(X_1) + log dR / dP^bbar, which cancels the unknown Z."""
-    log_weight = rollout["reward"] - rollout["cross"] - rollout["energy"]
-    return torch.softmax(log_weight, dim=0)  # [P] normalized weights, summing to one.
-
-
-def cross_entropy_loss(force: BiasForce, rollout: dict, weights: torch.Tensor) -> torch.Tensor:
-    """The weighted average of -log dP^{b_theta} / dR along the stored paths, the loss of 7.27.3."""
-    steps, paths = rollout["u_bar"].shape[0], rollout["u_bar"].shape[1]
-    dt = rollout["dt"]
-    clock = torch.arange(steps, device=weights.device).to(torch.get_default_dtype()) * dt
-    times = clock[:, None].expand(steps, paths)  # [T, P] the time of every stored state.
-    u = force(rollout["positions"].flatten(0, 1), rollout["velocities"].flatten(0, 1),
-              times.flatten()).view_as(rollout["u_bar"]) / (MASS * NOISE)  # [T, P, n, d].
-    drift = (u * (rollout["u_bar"] * dt + rollout["noise"])).sum(dim=(0, 2, 3))  # [P].
-    energy = 0.5 * (u ** 2).sum(dim=(0, 2, 3)) * dt  # [P] the control energy of each path.
-    return (weights * (energy - drift)).sum()  # Weighted by the importance weights.
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=100, lr=3e-3, width=64, samples=256)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    loud = not args.quiet
-    generator = torch.Generator().manual_seed(args.seed)
-
-    linear = lambda t: 0.3 * t  # The growth rate 0.3, so w^1_t rises at a constant rate.
-    late = lambda t: 0.3 * t ** 2  # The growth rate 0.6 t, which starts slowly.
-    near = {"linear": weighted_energy(2.0, lambda t: 1.0 - linear(t))
-            + weighted_energy(-2.0, linear)}
-    far = {"linear": weighted_energy(2.0, lambda t: 1.0 - linear(t)) + weighted_energy(-4.0, linear),
-           "late": weighted_energy(2.0, lambda t: 1.0 - late(t)) + weighted_energy(-4.0, late)}
-    half = float(branch_weights(linear, torch.tensor(0.5)))
-    if loud:
-        print("Two branches from the point 0, ending at +2 with weight 0.7 and at -2 with 0.3")
-        print(f"  weights at t = 1/2: {1 - half:.2f} and {half:.2f}; at t = 1: "
-              f"{1 - float(branch_weights(linear, torch.tensor(1.0))):.1f} and "
-              f"{float(branch_weights(linear, torch.tensor(1.0))):.1f}")
-        print(f"  weighted energies {weighted_energy(2.0, lambda t: 1.0 - linear(t)):.1f} and "
-              f"{weighted_energy(-2.0, linear):.1f}, total {near['linear']:.1f}")
-        print(f"  with the second target at -4: linear schedule {far['linear']:.1f}, "
-              f"late schedule {far['late']:.1f}, saving {far['linear'] - far['late']:.1f}")
-
-    bent = interpolant_energy(0.0, 2.0, 1.0)
-    straight = interpolant_energy(0.0, 2.0, 0.0)
-    grid = torch.tensor([0.0, 0.5, 1.0])
-    if loud:
-        print(f"\nThe interpolant with a correction of 1 passes through "
-              f"{[round(float(z), 2) for z in interpolant(0.0, 2.0, 1.0, grid)]}")
-        speed = [round(float(z), 0) for z in interpolant_velocity(0.0, 2.0, 1.0, grid)]
-        print(f"  its velocity is {speed} and its energy {bent:.4f}, against "
-              f"{straight:.1f} for the straight line")
-
-    reference_mass, inside, outside = 0.1, 1.0, 0.01  # The two-outcome example of Section 7.27.1.
-    normalizer = reference_mass * inside + (1.0 - reference_mass) * outside
-    tilted_mass = reference_mass * inside / normalizer
-    reward_average = tilted_mass * math.log(inside) + (1.0 - tilted_mass) * math.log(outside)
-    tilt_entropy = reward_average - math.log(normalizer)
-    if loud:
-        print(f"\nA reference reaching B one time in ten, tilted by pi_B = 1 inside and 0.01 outside")
-        print(f"  Z = {normalizer:.3f}, P*(B) = {tilted_mass:.3f}, "
-              f"a factor of {tilted_mass / reference_mass:.2f}")
-        print(f"  E[r] = {reward_average:.3f}, -log Z = {-math.log(normalizer):.3f}, "
-              f"KL(P* || R) = {tilt_entropy:.3f}, J(b*) = {tilt_entropy - reward_average:.3f}")
-
-    step_drift = (-0.3 + 0.5) * 0.01  # The physical force plus the bias, over one step.
-    step_friction = -1.0 * 0.8 * 0.01  # Minus the friction times the current velocity.
-    step_noise = NOISE * math.sqrt(0.01) * 0.4  # The noise scale times a normal draw of 0.4.
-    unit_control = 0.5 / NOISE  # The bias force measured in units of the noise.
-    plane = compose_bias(torch.tensor([[0.5]]), torch.tensor([[[0.3, -0.4]]]),
-                         torch.tensor([[[1.0, 0.0]]]))[0, 0]  # The example of Section 7.27.3.
-    log_weights = torch.tensor([-1.0, 0.0, 1.0])
-    softmax = torch.softmax(log_weights, dim=0)
-    if loud:
-        print(f"\nOne Euler-Maruyama step at dt = 0.01, with a noise scale of {NOISE:.4f}")
-        print(f"  dv = {step_drift:.4f} {step_friction:+.4f} {step_noise:+.4f} = "
-              f"{step_drift + step_friction + step_noise:.4f}, dr = {0.8 * 0.01:.4f}")
-        print(f"  the control in units of the noise is {unit_control:.4f}, contributing "
-              f"{0.5 * unit_control ** 2 * 0.01:.6f} per step and "
-              f"{0.5 * unit_control ** 2 * 0.01 * 100:.4f} over a hundred")
-        print(f"  b = {[round(float(z), 1) for z in plane]}, with alpha "
-              f"{float((plane * torch.tensor([1.0, 0.0])).sum()):.1f} along s_hat "
-              f"and length {float(plane.norm()):.3f}")
-        print(f"  three log-weights -1, 0, 1 normalize to {[round(float(z), 3) for z in softmax]}")
-
-    force = BiasForce(args.width).to(device)
-    reference_run = roll_out(None, args.samples, SIM_STEPS, device, generator)
-    before = roll_out(force, args.samples, SIM_STEPS, device, generator)
-
-    def loss_fn(step: int) -> torch.Tensor:
-        rollout = roll_out(force, PATHS, SIM_STEPS, device, generator)  # Simulate under b_bar.
-        return cross_entropy_loss(force, rollout, importance_weights(rollout))
-
-    if loud:
-        print(f"\ntraining the bias force, {count_parameters(force):,} parameters, "
-              f"{PATHS} paths of {SIM_STEPS} steps per optimizer step")
-    losses = train(force, loss_fn, steps=args.steps, lr=args.lr, quiet=not loud, log_every=50)
-    after = roll_out(force, args.samples, SIM_STEPS, device, generator)
-    effective = float(1.0 / (importance_weights(after) ** 2).sum() / args.samples)
-
-    report = {
-        "weight of branch 1 at t = 1/2": half,
-        "weighted energy, branch 0": weighted_energy(2.0, lambda t: 1.0 - linear(t)),
-        "weighted energy, branch 1": weighted_energy(-2.0, linear),
-        "total, both ends at distance 2": near["linear"],
-        "total with the far target, linear": far["linear"],
-        "total with the far target, late": far["late"],
-        "a branch drawn at t = 1/2": branch_sample(half, generator),
-        "interpolant energy, correction 1": bent,
-        "interpolant energy, straight": straight,
-        "Z of the two outcomes": normalizer,
-        "P*(B)": tilted_mass,
-        "the factor the tilt applies": tilted_mass / reference_mass,
-        "E_{P*}[r]": reward_average,
-        "KL(P* || R)": tilt_entropy,
-        "-log Z, the value of J(b*)": -math.log(normalizer),
-        "one velocity step": step_drift + step_friction + step_noise,
-        "control in units of the noise": unit_control,
-        "control energy over a hundred steps": 0.5 * unit_control ** 2 * 0.01 * 100,
-        "the composed force": [round(float(z), 2) for z in plane],
-        "its length": float(plane.norm()),
-        "softmax of -1, 0, 1": [round(float(z), 3) for z in softmax],
-        "mean reward, reference": float(reference_run["reward"].mean()),
-        "mean reward, before training": float(before["reward"].mean()),
-        "mean reward, after training": float(after["reward"].mean()),
-        "control energy after training": float(after["energy"].mean()),
-        "effective sample size, fraction": effective,
-        "force written by hand against autograd": force_check(after["positions"][-1]),
-        "final loss": sum(losses[-10:]) / len(losses[-10:]),
-    }
-    print_report("Branched and Entangled bridges, Chapter 7 Sections 7.26 to 7.27.3", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/csbm.py

@@ -1,254 +0,0 @@
-"""Categorical Schrodinger Bridge Matching: D-IMF on a path law small enough to write out.
-
-With N intermediate times a path is N + 2 states and the whole path law is a finite table, so both
-projections have one-line expressions. The reciprocal projection holds the endpoint coupling and
-redraws the interior from the reference, the Markov projection keeps the one-step conditionals and
-multiplies them into a chain, and the bridge is the only law in both sets. On two states and three
-transitions the sixteen paths fit in a list, so the alternation runs exactly.
-
-Chapter 7, Section 7.25 of the notes. Runs on S = 2, four time points, the uniform reference with
-alpha = 0.35, the source law (0.5, 0.5) and the requested law (0.2, 0.8). Prints the uniform and
-the ordered reference row, the exact bridge coupling, the conditionals of the first Markov
-projection, the gap after each round, and how that gap moves with alpha and with the path length.
-
-    python lecture_7/csbm.py              # 5 rounds of D-IMF, the five the chapter reports
-    python lecture_7/csbm.py --steps 1    # the single round the chapter writes out
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import torch
-
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-DTYPE = torch.float64  # The gap reaches 1e-11 in five rounds, well below float32.
-
-ALPHA = 0.35  # The probability of changing category in one step.
-TRANSITIONS = 3  # N + 1, so four time points and sixteen paths.
-PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE)  # The source law.
-PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE)  # The requested law at the last time.
-
-
-# %% 02. The reference over S categories
-#
-# The uniform reference holds the category with probability 1 - alpha and spreads the rest evenly
-# over the alternatives. For categories that have an order a discretized Gaussian centred on the
-# current one is available instead, under which a neighbour is likelier than a distant category.
-# Section 7.25.
-
-
-def uniform_reference(alpha: float, size: int) -> torch.Tensor:
-    """A_n(x, y) = 1 - alpha on the diagonal and alpha / (S - 1) off it. [S, S]."""
-    off = alpha / (size - 1)  # The share each alternative receives.
-    return torch.full((size, size), off, dtype=DTYPE) + torch.eye(size, dtype=DTYPE) * (
-        1.0 - alpha - off)
-
-
-def ordered_reference(width: float, size: int) -> torch.Tensor:
-    """Rows proportional to exp(-(y - x)^2 / 2 width^2), under which a nearby category is likelier."""
-    offset = torch.arange(size, dtype=DTYPE)[None, :] - torch.arange(size, dtype=DTYPE)[:, None]
-    weights = torch.exp(-offset ** 2 / (2.0 * width ** 2))  # [S, S] before normalizing.
-    return weights / weights.sum(dim=1, keepdim=True)
-
-
-def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int = 400
-             ) -> torch.Tensor:
-    """Scale diag(p_0) A^N to the two marginals, which gives the exact bridge coupling."""
-    f = torch.ones_like(start)
-    g = torch.ones_like(end)
-    for _ in range(sweeps):
-        g = end / (matrix.T @ f)  # Column sums become exactly p_1.
-        f = start / (matrix @ g)  # Row sums become exactly p_0.
-    return f[:, None] * matrix * g[None, :]
-
-
-# %% 03. The two projections in discrete time
-#
-# A path law is a tensor with one axis per time. The reciprocal projection multiplies the endpoint
-# coupling by the reference bridge weight of every path, and the Markov projection reads off the
-# one-step conditionals and multiplies them back into a chain. Section 7.25.1.
-
-
-def reference_joint(matrix: torch.Tensor, steps: int) -> torch.Tensor:
-    """The product of the transition probabilities along a path, as a tensor with one axis per time."""
-    joint = matrix.clone()  # [x_0, x_1].
-    for _ in range(steps - 1):
-        joint = joint.unsqueeze(-1) * matrix  # Append one more transition on a new axis.
-    return joint  # [S] * (steps + 1).
-
-
-def bridge_table(matrix: torch.Tensor, steps: int) -> torch.Tensor:
-    """R(x_interior | x_0, x_N), the reference bridge weight of every path given its two ends."""
-    whole = torch.linalg.matrix_power(matrix, steps)  # [S, S] the reference over the interval.
-    size = len(matrix)
-    shape = [size] + [1] * (steps - 1) + [size]  # Broadcast over the interior axes.
-    return reference_joint(matrix, steps) / whole.reshape(shape)
-
-
-def reciprocal_projection(coupling: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor:
-    """proj_R(P) = P(x_0, x_N) R(x_in | x_0, x_N), the endpoint coupling with reference bridges."""
-    size = len(matrix)
-    shape = [size] + [1] * (steps - 1) + [size]
-    return coupling.reshape(shape) * bridge_table(matrix, steps)
-
-
-def markov_projection(law: torch.Tensor) -> tuple[list[torch.Tensor], torch.Tensor]:
-    """proj_M(P) = P(x_0) prod_n P(x_{t_n} | x_{t_{n-1}}), as the conditionals and the coupling."""
-    conditionals = []
-    for step in range(law.dim() - 1):
-        axes = [d for d in range(law.dim()) if d not in (step, step + 1)]
-        pair = law.sum(dim=axes)  # [S, S] the joint of two consecutive states.
-        conditionals.append(pair / pair.sum(dim=1, keepdim=True))
-    joint = torch.diag(law.sum(dim=list(range(1, law.dim()))))  # [S, S] the law at time 0.
-    for matrix in conditionals:
-        joint = joint @ matrix
-    return conditionals, joint
-
-
-def transition_loss(true: torch.Tensor, network: torch.Tensor) -> float:
-    """KL(P(x_t | x_s) || m_theta(x_t | x_s)) on one transition out of one state, in nats."""
-    return float((true * (true / network).log()).sum())
-
-
-# %% 04. D-IMF on sixteen paths
-#
-# Alternating the two projections is iterative Markovian fitting with discrete time. The starting
-# law is the independent coupling filled in with reference bridges, so both requested marginals are
-# already exact and only the joint law over paths moves. Section 7.25.2.
-
-
-def d_imf(coupling: torch.Tensor, matrix: torch.Tensor, steps: int, rounds: int,
-          target: torch.Tensor) -> tuple[list[float], list[torch.Tensor]]:
-    """Run the alternation, returning the gap after each round and the first round's conditionals."""
-    gaps: list[float] = []
-    first: list[torch.Tensor] = []
-    for index in range(rounds):
-        law = reciprocal_projection(coupling, matrix, steps)  # Reference bridges in the interior.
-        conditionals, coupling = markov_projection(law)  # The chain and its endpoint coupling.
-        if index == 0:
-            first = conditionals
-        gaps.append(float((reciprocal_projection(coupling, matrix, steps) - target).abs().sum()))
-    return gaps, first
-
-
-def starting_gap(matrix: torch.Tensor, steps: int, start: torch.Tensor,
-                 end: torch.Tensor) -> tuple[float, torch.Tensor, torch.Tensor]:
-    """The gap of the independent coupling from the bridge, with both path laws."""
-    whole = torch.linalg.matrix_power(matrix, steps)
-    plan = sinkhorn(torch.diag(start) @ whole, start, end)
-    target = reciprocal_projection(plan, matrix, steps)
-    initial = reciprocal_projection(start[:, None] * end[None, :], matrix, steps)
-    return float((initial - target).abs().sum()), initial, target
-
-
-# %% 05. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=5)  # --steps is the number of D-IMF rounds.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    rounds = max(args.steps, 1)
-
-    A = uniform_reference(ALPHA, len(PI_0))
-    whole = torch.linalg.matrix_power(A, TRANSITIONS)
-    ordered = ordered_reference(1.0, 5)  # Five ordered categories, a width of one category.
-    if loud:
-        print(f"S = 2, {TRANSITIONS + 1} time points, alpha = {ALPHA}, "
-              f"{len(PI_0) ** (TRANSITIONS + 1)} paths")
-        print(f"  A = {[[round(float(z), 2) for z in row] for row in A]}, "
-              f"A^{TRANSITIONS} = {[[round(float(z), 4) for z in row] for row in whole]}")
-        print(f"  the ordered reference out of the middle of five categories: "
-              f"{[round(float(z), 4) for z in ordered[2]]}")
-
-    plan = sinkhorn(torch.diag(PI_0) @ whole, PI_0, PI_1)
-    gap_0, initial, target = starting_gap(A, TRANSITIONS, PI_0, PI_1)
-    interior = [initial.sum(dim=[d for d in range(TRANSITIONS + 1) if d != k]) for k in (1, 2)]
-    independent = PI_0[:, None] * PI_1[None, :]  # Where the alternation starts.
-    if loud:
-        print(f"\npi* = {[[round(float(z), 5) for z in row] for row in plan]}, against the "
-              f"independent coupling "
-              f"{[[round(float(z), 1) for z in row] for row in independent]}")
-        print(f"  the path 1,1,1,1 has reference bridge weight "
-              f"{float(bridge_table(A, TRANSITIONS)[0, 0, 0, 0]):.4f}, so its probability is "
-              f"{float(initial[0, 0, 0, 0]):.4f} at the start and {float(target[0, 0, 0, 0]):.4f} "
-              f"under the bridge")
-        print(f"  the path 2,2,2,2 has {float(initial[1, 1, 1, 1]):.4f} and "
-              f"{float(target[1, 1, 1, 1]):.4f}")
-        print(f"  the two laws agree at the interior times, "
-              f"{[[round(float(z), 4) for z in m] for m in interior]}, and differ by "
-              f"{gap_0:.4f} over the paths")
-
-    gaps, conditionals = d_imf(independent, A, TRANSITIONS, rounds, target)
-    # Eight rounds reach machine precision, past which a ratio of two gaps measures rounding, so
-    # the report quotes the first one.
-    ratios = [gaps[k] / gaps[k + 1] for k in range(len(gaps) - 1) if gaps[k + 1] > 1e-14]
-    per_round = ratios[0] if ratios else gap_0 / gaps[0]  # One round alone gives only the first.
-    network = torch.tensor([0.25, 0.75], dtype=DTYPE)  # A network that is wrong by a little.
-    one_term = transition_loss(conditionals[-1][0], network)
-    if loud:
-        print(f"\nthe {len(conditionals)} conditionals of the first Markov projection")
-        for matrix in conditionals:
-            print(f"    {[[round(float(z), 4) for z in row] for row in matrix]}")
-        print(f"  gap {gap_0:.4f} at the start, then {['%.1e' % x for x in gaps]}")
-        print(f"  ratios {[round(r, 1) for r in ratios]}")
-        print(f"  the true conditional {[round(float(z), 4) for z in conditionals[-1][0]]} against "
-              f"a network at {[float(z) for z in network]} costs {one_term:.4f} nats")
-
-    # A reference that mixes well and a longer chain both leave the bridge less to do, which the
-    # starting gap records and the rate of convergence follows. Two rounds separate the two ratios
-    # from the float64 floor, which the sixty-four-path column reaches after three.
-    mixing = {}
-    for alpha in (0.2, ALPHA, 0.4):
-        matrix = uniform_reference(alpha, len(PI_0))
-        start, _, goal = starting_gap(matrix, TRANSITIONS, PI_0, PI_1)
-        later, _ = d_imf(independent, matrix, TRANSITIONS, 2, goal)
-        mixing[alpha] = (start, later[0] / later[1])
-    lengths = {}
-    for steps in (2, TRANSITIONS, 5):
-        start, _, goal = starting_gap(A, steps, PI_0, PI_1)
-        later, _ = d_imf(independent, A, steps, 2, goal)
-        lengths[steps] = (start, later[0] / later[1])
-    if loud:
-        print("\nalpha   starting gap   gap divided per round")
-        for alpha, (start, ratio) in mixing.items():
-            print(f"  {alpha:<6g} {start:<14.5f} {ratio:.0f}")
-        print("\npaths   starting gap   gap divided per round")
-        for steps, (start, ratio) in lengths.items():
-            print(f"  {len(PI_0) ** (steps + 1):<6d} {start:<14.5f} {ratio:.0f}")
-
-    report = {
-        "A^3": [round(float(z), 4) for z in whole.flatten()],
-        "the ordered reference row": [round(float(z), 4) for z in ordered[2]],
-        "pi*_11": float(plan[0, 0]),
-        "pi*_12": float(plan[0, 1]),
-        "reference bridge weight of 1,1,1,1": float(bridge_table(A, TRANSITIONS)[0, 0, 0, 0]),
-        "P^0(1,1,1,1)": float(initial[0, 0, 0, 0]),
-        "P*(1,1,1,1)": float(target[0, 0, 0, 0]),
-        "P^0(2,2,2,2)": float(initial[1, 1, 1, 1]),
-        "P*(2,2,2,2)": float(target[1, 1, 1, 1]),
-        "marginal at the first interior time": [round(float(z), 4) for z in interior[0]],
-        "marginal at the second": [round(float(z), 4) for z in interior[1]],
-        "gap at the start": gap_0,
-        "gap after one round": gaps[0],
-        "gap after the last round": gaps[-1],
-        "rounds run": rounds,
-        "gap divided per round": per_round,
-        "first conditional, out of state 1": [round(float(z), 4) for z in conditionals[0][0]],
-        "last conditional, out of state 1": [round(float(z), 4) for z in conditionals[-1][0]],
-        "one transition against a wrong network": one_term,
-        "starting gap at alpha = 0.2": mixing[0.2][0],
-        "rounds divide the gap by, at alpha = 0.2": mixing[0.2][1],
-        "rounds divide the gap by, at alpha = 0.4": mixing[0.4][1],
-        "starting gap on eight paths": lengths[2][0],
-        "starting gap on sixty-four paths": lengths[5][0],
-        "rounds divide the gap by, on eight paths": lengths[2][1],
-        "rounds divide the gap by, on sixty-four paths": lengths[5][1],
-    }
-    print_report("Categorical Schrodinger Bridge Matching, Chapter 7 Section 7.25", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/ddsbm.py

@@ -1,243 +0,0 @@
-"""Discrete Diffusion Schrodinger Bridge Matching: one categorical feature of a molecular graph.
-
-A molecule is a graph with a categorical feature on every atom and on every pair of atoms, and the
-reference resamples each feature independently, so one feature with S categories specifies the
-whole construction. Scaling the reference kernel to the two laws gives the bridge coupling for that
-feature, and pinning a path at a known terminal graph gives the rate a network regresses onto.
-Both stay usable on a state space of 10^27 graphs because the reference is a product over the
-features, so every sum over graphs becomes a sum over the categories of one of them.
-
-Chapter 7, Section 7.24 of the notes. Runs on one bond with three types, no bond, single and
-double, at the survival probability r = 0.4. Prints the kernel and its generator, the bridge
-coupling and the mass it leaves unchanged, the pinned rates and the posterior average a network
-learns, and the counting over the forty-five features of a nine-atom graph.
-
-    python lecture_7/ddsbm.py              # 200 Sinkhorn sweeps on the bond coupling
-    python lecture_7/ddsbm.py --steps 1    # the first sweep alone, computed by hand
-"""
-
-# %% 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  # The chapter quotes four figures of the coupling and of the pinned rates.
-
-PRIOR = torch.full((3,), 1.0 / 3.0, dtype=DTYPE)  # m, uniform over the three bond types.
-BOND_0 = torch.tensor([0.7, 0.2, 0.1], dtype=DTYPE)  # pi_0, mostly no bond.
-BOND_1 = torch.tensor([0.1, 0.2, 0.7], dtype=DTYPE)  # pi_1, mostly double.
-SURVIVAL = 0.4  # r = alpha_bar(1) / alpha_bar(0) over the whole unit interval.
-ATOMS = 9  # Heavy atoms in the molecule the chapter counts.
-ATOM_TYPES, BOND_TYPES = 5, 4  # Categories of a node feature and of a bond feature.
-
-
-# %% 02. The reference process on one feature
-#
-# The kernel keeps the category with probability r and redraws it from the prior m otherwise. Over
-# a short step r = 1 - h beta(t), so the rate out of x toward a different category y is the
-# resampling rate beta(t) times the chance m(y) that the redraw lands on y. Section 7.24.1.
-
-
-def replacement_kernel(r: float, prior: torch.Tensor) -> torch.Tensor:
-    """K_st(x, y) = r 1_{x=y} + (1 - r) m(y), the uniform-replacement kernel with a general prior."""
-    return r * torch.eye(len(prior), dtype=DTYPE) + (1.0 - r) * prior[None, :]  # [S, S].
-
-
-def replacement_generator(beta: float, prior: torch.Tensor) -> torch.Tensor:
-    """Q_t(x, y) = beta(t) m(y) off the diagonal, with the diagonal -beta(t) (1 - m(x))."""
-    rates = beta * prior[None, :].repeat(len(prior), 1)  # [S, S].
-    rates = rates - torch.diag(rates.diagonal())  # Clear the diagonal before closing the row.
-    return rates - torch.diag(rates.sum(dim=1))
-
-
-def resampling_rate(r: float) -> float:
-    """beta = -log alpha_bar(1) for the schedule alpha_bar(t) = r^t, which is constant in time."""
-    return -math.log(r)
-
-
-# %% 03. A bridge for a single bond
-#
-# Scaling diag(pi_0) K_01 to the two marginals gives the coupling. The mass it leaves on the
-# diagonal sits between the mass an independent coupling happens to leave and the largest any
-# coupling can leave, and r fixes where between the two it lands. Section 7.24.2.
-
-
-def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int
-             ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[float]]:
-    """Scale a positive matrix to the two marginals, returning f, g, the plan and the column errors."""
-    f = torch.ones_like(start)
-    g = torch.ones_like(end)
-    errors: list[float] = []
-    for _ in range(sweeps):
-        g = end / (matrix.T @ f)  # Column sums become exactly pi_1.
-        f = start / (matrix @ g)  # Row sums become exactly pi_0, and the columns move again.
-        errors.append(float(((f[:, None] * matrix * g[None, :]).sum(dim=0) - end).abs().sum()))
-    return f, g, f[:, None] * matrix * g[None, :], errors
-
-
-def held_mass(plan: torch.Tensor) -> float:
-    """The trace of the plan, the fraction of bonds whose type the bridge does not change."""
-    return float(plan.diagonal().sum())
-
-
-# %% 04. Training with the terminal graph fixed
-#
-# Pinning the endpoint is an h-transform with the reference probability of arriving at x_1 in the
-# place of the potential, so a move onto the terminal type has its rate multiplied and a move off
-# it divided. Several terminal types can occur at one state, and the minimizer of the rate error
-# over a single number is their posterior average. Section 7.24.3.
-
-
-def pinned_rate(generator: torch.Tensor, kernel: torch.Tensor, x: int, y: int, x1: int) -> float:
-    """Q_t(x, y; x_1) = Q_t(x, y) K_t1(y, x_1) / K_t1(x, x_1), the rate of the bridge ending at x_1."""
-    return float(generator[x, y] * kernel[y, x1] / kernel[x, x1])
-
-
-def half_survival_time(r: float) -> float:
-    """The time where alpha_bar(1) / alpha_bar(t) = r^{1-t} equals 1/2, on alpha_bar(t) = r^t."""
-    return 1.0 - math.log(0.5) / math.log(r)  # At r = 0.4 this is log 0.8 / log 0.4 = 0.2435.
-
-
-def rate_error(controlled: float, network: float) -> float:
-    """Q^pin log(Q^pin / Q^theta) - Q^pin + Q^theta, the Bregman divergence of u log u on one edge."""
-    return controlled * math.log(controlled / network) - controlled + network
-
-
-def posterior_loss(rates: list[float], weights: list[float], network: float) -> float:
-    """The expected rate error at a state whose terminal type is uncertain, over the posterior."""
-    return sum(w * rate_error(q, network) for q, w in zip(rates, weights))
-
-
-# %% 05. The factorization over a graph
-#
-# The reference applies the same kernel independently at every node and every pair of nodes, so the
-# generator changes one feature at a time, the total jump rate out of a graph is the sum of the
-# per-feature rates, and in the pinned rate every factor cancels but the one of the feature that
-# moves. Section 7.24.4.
-
-
-def feature_count(atoms: int) -> tuple[int, int, int]:
-    """The node features, the bond features and their total, for a molecule with this many atoms."""
-    bonds = atoms * (atoms - 1) // 2  # One feature per unordered pair.
-    return atoms, bonds, atoms + bonds
-
-
-def state_space_size(atoms: int, atom_types: int, bond_types: int) -> float:
-    """atom_types^nodes times bond_types^bonds, the number of graphs the sums would run over."""
-    nodes, bonds, _ = feature_count(atoms)
-    return float(atom_types ** nodes) * float(bond_types ** bonds)
-
-
-def total_jump_rate(features: int, beta: float, categories: int) -> float:
-    """The sum of the per-feature rates out of a graph, which fixes the step size of the sampler."""
-    return features * beta * (1.0 - 1.0 / categories)  # A redraw lands elsewhere with 1 - m(x).
-
-
-# %% 06. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=200)  # --steps is the number of Sinkhorn sweeps on the bond coupling.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    sweeps = max(args.steps, 1)
-
-    kernel = replacement_kernel(SURVIVAL, PRIOR)
-    beta = resampling_rate(SURVIVAL)
-    rates = replacement_generator(beta, PRIOR)
-    reference = torch.diag(BOND_0) @ kernel
-    if loud:
-        print(f"One bond with three types at r = {SURVIVAL}, beta = {beta:.4f}, "
-              f"off-diagonal rate {float(rates[0, 1]):.4f}")
-        print(f"  K_01 = {[[round(float(z), 1) for z in row] for row in kernel]}")
-        print(f"  diag(pi_0) K_01 has column sums "
-              f"{[round(float(z), 2) for z in reference.sum(dim=0)]}, so the first g is "
-              f"{[round(float(z), 4) for z in BOND_1 / reference.sum(dim=0)]}")
-
-    f, g, plan, errors = sinkhorn(reference, BOND_0, BOND_1, sweeps)
-    independent = float((BOND_0 * BOND_1).sum())
-    largest = float(torch.minimum(BOND_0, BOND_1).sum())
-    held = {r: held_mass(sinkhorn(torch.diag(BOND_0) @ replacement_kernel(r, PRIOR),
-                                  BOND_0, BOND_1, 200)[2])
-            for r in (0.2, 0.4, 0.6, 0.8, 0.99)}
-    if loud:
-        print(f"\nAfter {sweeps} sweeps, column error {errors[-1]:.2e}")
-        print(f"  f = {[round(float(z), 4) for z in f]}, "
-              f"g = {[round(float(z), 4) for z in g]}")
-        print(f"  pi* = {[[round(float(z), 3) for z in row] for row in plan]}")
-        print(f"  mass that stays: bridge {held_mass(plan):.4f}, independent {independent:.2f}, "
-              f"largest possible {largest:.1f}")
-        print(f"  diagonal mass at r = 0.2, 0.4, 0.6, 0.8, 0.99: "
-              f"{[round(held[r], 4) for r in (0.2, 0.4, 0.6, 0.8, 0.99)]}")
-
-    pinned_time = half_survival_time(SURVIVAL)
-    pinned_kernel = replacement_kernel(0.5, PRIOR)
-    toward = pinned_rate(rates, pinned_kernel, 0, 2, 2)  # A move onto the terminal type.
-    away = pinned_rate(rates, pinned_kernel, 2, 0, 2)  # A move off it.
-    between = pinned_rate(rates, pinned_kernel, 0, 1, 2)  # Between two non-terminal types.
-    one_edge = rate_error(toward, 1.0)  # The term one edge contributes at a network rate of 1.
-    weights = [0.6, 0.4]  # The posterior over the two terminal types at this state.
-    posterior = weights[0] * toward + weights[1] * between
-    losses = {q: posterior_loss([toward, between], weights, q) for q in (posterior, 1.0, 0.5)}
-    if loud:
-        print(f"\nAt t = {pinned_time:.4f} the survival ratio to time 1 is 1/2, so K_t1 is "
-              f"{float(pinned_kernel[0, 0]):.4f} on the diagonal and "
-              f"{float(pinned_kernel[0, 1]):.4f} off it")
-        print(f"  pinned rates: {toward:.4f} toward the terminal type, {away:.4f} away from it "
-              f"and {between:.4f} between two others, a factor of {toward / away:.0f}")
-        print(f"  one edge at a network rate of 1 contributes {one_edge:.4f}")
-        print(f"  with two terminal types the network learns {posterior:.4f}, at a loss of "
-              f"{losses[posterior]:.4f}, against {losses[1.0]:.4f} at 1 and {losses[0.5]:.4f} at 0.5")
-
-    nodes, bonds, features = feature_count(ATOMS)
-    graphs = state_space_size(ATOMS, ATOM_TYPES, BOND_TYPES)
-    jump_rate = total_jump_rate(bonds, beta, BOND_TYPES)
-    if loud:
-        print(f"\n{ATOMS} heavy atoms: {nodes} node features and {bonds} bond features, "
-              f"{features} in all")
-        print(f"  {ATOM_TYPES}^{nodes} x {BOND_TYPES}^{bonds} = {graphs:.1e} graphs, "
-              f"which no sum reaches")
-        print(f"  the total jump rate out of a graph is {bonds} x {beta:.4f} x "
-              f"(1 - 1/{BOND_TYPES}) = {jump_rate:.2f} per unit time")
-
-    report = {
-        "beta of alpha_bar(t) = 0.4^t": beta,
-        "off-diagonal reference rate": float(rates[0, 1]),
-        "K_01 on the diagonal": float(kernel[0, 0]),
-        "K_01 off the diagonal": float(kernel[0, 1]),
-        "rows of the generator sum to": float(rates.sum(dim=1).abs().max()),
-        "column sums of diag(pi_0) K_01": [round(float(z), 2) for z in reference.sum(dim=0)],
-        "the first g": [round(float(z), 4) for z in BOND_1 / reference.sum(dim=0)],
-        "sweeps run": sweeps,
-        "column error after the last sweep": errors[-1],
-        "f": [round(float(z), 4) for z in f],
-        "g": [round(float(z), 4) for z in g],
-        "bond plan, no bond to double": float(plan[0, 2]),
-        "bond mass that stays": held_mass(plan),
-        "bond mass, independent coupling": independent,
-        "bond mass, largest possible": largest,
-        "diagonal mass at r = 0.2": held[0.2],
-        "diagonal mass at r = 0.99": held[0.99],
-        "time where the survival ratio is 1/2": pinned_time,
-        "pinned rate toward the terminal": toward,
-        "pinned rate away from it": away,
-        "pinned rate between two others": between,
-        "one edge at a network rate of 1": one_edge,
-        "posterior average of the rates": posterior,
-        "loss at the posterior average": losses[posterior],
-        "loss at 1": losses[1.0],
-        "loss at 0.5": losses[0.5],
-        "node, bond and total features": [nodes, bonds, features],
-        "graphs on nine heavy atoms": f"{graphs:.1e}",
-        "total jump rate out of a graph": jump_rate,
-    }
-    print_report("Discrete Diffusion Schrodinger Bridge Matching, Chapter 7 Section 7.24", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/discrete_bridge.py

@@ -1,267 +0,0 @@
-"""The Schrodinger bridge over a finite set: rates in place of drifts.
-
-On a finite state space there is no small displacement, so the control is a reweighting of jump
-rates and relative entropy on paths is a time integral of a Bregman divergence between two rate
-matrices. Everything else carries over. The chain rule splits the problem into an endpoint term and
-a bridge term, the two potentials are the Sinkhorn scalings of the reference endpoint law, and
-under Doob's h-transform each reference rate is multiplied by the ratio of a backward potential
-at the destination to its value at the current state.
-
-Chapter 7, Sections 7.21 to 7.23.3 of the notes. Runs the two-state continuous-time example with
-the generator Q, the source law (0.5, 0.5) and the requested terminal law (0.2, 0.8). Prints the
-bridge rates at three times, the reverse rates, the Euler sampler's error as the step shrinks,
-and the cost computed on the endpoints and on the paths.
-
-    python lecture_7/discrete_bridge.py              # 10 Euler steps, then 20, 40 and 100
-    python lecture_7/discrete_bridge.py --steps 20   # a finer sampler, 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  # Matrix exponentials and a time integral, both in double precision.
-
-Q = torch.tensor([[-1.0, 1.0], [1.5, -1.5]], dtype=DTYPE)  # The reference generator.
-PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE)  # The source law.
-PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE)  # The requested terminal law.
-STATIONARY = torch.tensor([0.6, 0.4], dtype=DTYPE)  # Proportional to (1.5, 1), the rows of Q.
-GRID = 1001  # Points in the Simpson rule for the time integral of the rate error.
-
-
-# %% 02. Relative entropy between two jump processes
-#
-# Collecting the one-step relative entropies and letting the step go to zero leaves a time integral
-# of the Bregman divergence of u log u between the two rates, which is nonnegative and vanishes
-# only when they agree. The divergence is not symmetric in the ratio. Section 7.21.1.
-
-
-def rate_error(controlled: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
-    """Q^theta log(Q^theta / Q) - Q^theta + Q, the Bregman divergence of u log u, entry by entry."""
-    safe = controlled.clamp(min=1e-300)  # A zero rate contributes the reference rate alone.
-    return controlled * (safe / reference).log() - controlled + reference
-
-
-def scaling_cost(c: float) -> float:
-    """c log c - c + 1, the cost per unit of reference rate of scaling every rate by c."""
-    return c * math.log(c) - c + 1.0
-
-
-# %% 03. The Schrodinger problem for a chain
-#
-# The endpoint law of a chain with a time-independent generator is diag(pi_0^ref) e^Q, and the two
-# potentials are the Sinkhorn scalings of that matrix to the two marginals. Section 7.21.2.
-
-
-def endpoint_law(generator: torch.Tensor, start: torch.Tensor) -> torch.Tensor:
-    """R_01 = diag(pi_0) e^Q, the reference probability of starting at x and ending at y."""
-    return torch.diag(start) @ torch.matrix_exp(generator)  # [S, S].
-
-
-def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int = 400
-             ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Scale a positive matrix to the two marginals, returning f, g and the plan f R_01 g."""
-    f = torch.ones_like(start)
-    g = torch.ones_like(end)
-    for _ in range(sweeps):
-        g = end / (matrix.T @ f)  # Column sums become exactly pi_1.
-        f = start / (matrix @ g)  # Row sums become exactly pi_0.
-    return f, g, f[:, None] * matrix * g[None, :]
-
-
-def relative_entropy(p: torch.Tensor, q: torch.Tensor) -> float:
-    """Sum over entries of p log(p / q), for two positive matrices or vectors."""
-    mask = p > 0
-    return float((p[mask] * (p[mask] / q[mask]).log()).sum())
-
-
-# %% 04. Rates reweighted by an h-transform
-#
-# The backward potential is g averaged over the reference from the current time to time 1, and each
-# off-diagonal rate is multiplied by the ratio of that potential at the destination to its value at
-# the current state. The diagonal is then set so that the row sums to zero. Section 7.22.
-
-
-def backward_potential(generator: torch.Tensor, g: torch.Tensor, t: float) -> torch.Tensor:
-    """phi_t = e^{Q(1-t)} g, the reference average of g at the final state seen from time t."""
-    return torch.matrix_exp(generator * (1.0 - t)) @ g  # [S].
-
-
-def tilted_generator(generator: torch.Tensor, phi: torch.Tensor) -> torch.Tensor:
-    """Q*(x, y) = Q(x, y) phi(y) / phi(x) off the diagonal, with the diagonal closing each row."""
-    tilted = generator * phi[None, :] / phi[:, None]  # [S, S] every entry reweighted.
-    tilted.fill_diagonal_(0.0)  # Discard the reweighted diagonal, which has no meaning.
-    return tilted - torch.diag(tilted.sum(dim=1))  # Each row now sums to zero.
-
-
-def bridge_marginal(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor,
-                    g: torch.Tensor, t: float) -> torch.Tensor:
-    """p*_t = psi_t phi_t, the forward potential carried under the reference times the backward one."""
-    psi = (start * f) @ torch.matrix_exp(generator * t)  # [S] the forward factor.
-    return psi * backward_potential(generator, g, t)  # [S] their entrywise product.
-
-
-# %% 05. Sampling the bridge with Euler steps
-#
-# Over a step the chain moves to a different state with probability dt Q*(x, y) and stays with the
-# remaining probability, which is the Euler sampler of Chapter 4 applied to a generator that
-# changes with time. Section 7.23.1.
-
-
-def euler_marginal(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor,
-                   g: torch.Tensor, steps: int) -> tuple[torch.Tensor, list[float]]:
-    """Propagate p <- p (I + dt Q*_t) and record the probability of leaving state 1 at each step."""
-    dt = 1.0 / steps
-    p = start.clone()
-    leaving = []
-    identity = torch.eye(len(start), dtype=DTYPE)
-    for step in range(steps):
-        tilted = tilted_generator(generator, backward_potential(generator, g, step * dt))
-        leaving.append(float(tilted[0, 1] * dt))  # The jump probability out of state 1.
-        p = p @ (identity + dt * tilted)  # One Euler step of the master equation.
-    return p, leaving
-
-
-# %% 06. The same bridge run backwards
-#
-# Matching the probability flux on an edge in both directions fixes the reverse rate, which is the
-# forward rate times the ratio of the marginal at the two ends of the edge. Section 7.23.2.
-
-
-def reverse_rates(tilted: torch.Tensor, marginal: torch.Tensor) -> torch.Tensor:
-    """Q_hat(y, x) = Q(x, y) p_t(x) / p_t(y), the rates under which the flux on an edge matches."""
-    return (tilted * marginal[:, None] / marginal[None, :]).T  # [S, S], reversed in both indices.
-
-
-# %% 07. The cost of the bridge, computed twice
-#
-# On the endpoints the cost is the relative entropy of the plan from the reference endpoint law. On
-# the paths it is the time integral of the rate error, weighted by the marginal of the bridge at
-# that time. The two agree. Section 7.23.3.
-
-
-def path_cost(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor, g: torch.Tensor,
-              points: int = GRID) -> tuple[float, float]:
-    """Integrate the weighted rate error over [0, 1] by Simpson's rule.
-
-    Returns the integral and the value of the integrand at t = 1/2, the one the chapter prints.
-    """
-    off = ~torch.eye(len(generator), dtype=torch.bool)  # [S, S] true away from the diagonal.
-    times = torch.linspace(0.0, 1.0, points, dtype=DTYPE)
-    values = []
-    for t in times.tolist():
-        tilted = tilted_generator(generator, backward_potential(generator, g, t))
-        marginal = bridge_marginal(generator, start, f, g, t)
-        error = torch.zeros_like(tilted)
-        error[off] = rate_error(tilted[off], generator[off])  # Only the edges carry a cost.
-        values.append(float((marginal * error.sum(dim=1)).sum()))
-    weights = torch.ones(points, dtype=DTYPE)
-    weights[1:-1:2] = 4.0  # Simpson's rule on an even number of intervals.
-    weights[2:-1:2] = 2.0
-    integral = float((weights * torch.tensor(values, dtype=DTYPE)).sum() / (3.0 * (points - 1)))
-    return integral, values[(points - 1) // 2]
-
-
-# %% 08. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=10)  # --steps is the coarsest number of Euler steps in the table.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    coarse = max(args.steps, 2)
-
-    if loud:
-        print("Scaling every reference rate by c costs c log c - c + 1 per unit of rate")
-        print(f"  halving {scaling_cost(0.5):.4f}, doubling {scaling_cost(2.0):.4f}")
-
-    R01 = endpoint_law(Q, PI_0)
-    f, g, plan = sinkhorn(R01, PI_0, PI_1)
-    phi_half = backward_potential(Q, g, 0.5)
-    tilted_half = tilted_generator(Q, phi_half)
-    tilted_zero = tilted_generator(Q, backward_potential(Q, g, 0.0))
-    tilted_one = tilted_generator(Q, g)
-    marginal_half = bridge_marginal(Q, PI_0, f, g, 0.5)
-    marginal_one = bridge_marginal(Q, PI_0, f, g, 1.0)
-    reverse = reverse_rates(tilted_half, marginal_half)
-    if loud:
-        print(f"\ne^Q = {[[round(float(z), 4) for z in row] for row in torch.matrix_exp(Q)]}")
-        print(f"R_01 = {[[round(float(z), 4) for z in row] for row in R01]}, "
-              f"column sums {[round(float(z), 4) for z in R01.sum(0)]}")
-        print(f"  f = {[round(float(z), 4) for z in f]}, g = {[round(float(z), 4) for z in g]}, "
-              f"g ratio {float(g[1] / g[0]):.3f}")
-        print(f"  pi* = {[[round(float(z), 4) for z in row] for row in plan]}")
-        print(f"  phi_1/2 = {[round(float(z), 4) for z in phi_half]}, "
-              f"ratio {float(phi_half[1] / phi_half[0]):.2f}")
-        print(f"  Q*_1/2 = {[[round(float(z), 4) for z in row] for row in tilted_half]}")
-        print(f"  rates out of the two states: at t = 0 "
-              f"{float(tilted_zero[0, 1]):.4f} and {float(tilted_zero[1, 0]):.4f}, "
-              f"at t = 1 {float(tilted_one[0, 1]):.4f} and {float(tilted_one[1, 0]):.4f}")
-        print(f"  marginal at t = 1/2 {[round(float(z), 4) for z in marginal_half]}, "
-              f"at t = 1 {[round(float(z), 4) for z in marginal_one]}")
-        print(f"  reverse rates at t = 1/2: {float(reverse[1, 0]):.4f} from 2 back to 1, "
-              f"{float(reverse[0, 1]):.4f} from 1 back to 2")
-        print(f"  the flux on the edge: {float(marginal_half[0] * tilted_half[0, 1]):.4f} "
-              f"against {float(marginal_half[1] * reverse[1, 0]):.4f}")
-
-    sampler = {}
-    for steps in (coarse, 2 * coarse, 4 * coarse, 10 * coarse):
-        terminal, leaving = euler_marginal(Q, PI_0, f, g, steps)
-        sampler[steps] = (terminal, float((terminal - PI_1).abs().sum() / 2), leaving)
-    if loud:
-        print("\nsteps   terminal law                 error    first and last jump probability")
-        for steps, (terminal, error, leaving) in sampler.items():
-            print(f"  {steps:<5d} ({float(terminal[0]):.4f}, {float(terminal[1]):.4f})"
-                  f"           {error:.4f}   {leaving[0]:.4f}, {leaving[-1]:.4f}")
-
-    endpoint_cost = relative_entropy(plan, R01)
-    integral, at_half = path_cost(Q, PI_0, f, g)
-    speeds = {}
-    for lam in (0.5, 1.0, 2.0, 4.0):
-        scaled_law = endpoint_law(lam * Q, PI_0)
-        _, _, scaled_plan = sinkhorn(scaled_law, PI_0, PI_1)
-        speeds[lam] = relative_entropy(scaled_plan, scaled_law)
-    limit = relative_entropy(PI_1, STATIONARY)
-    if loud:
-        print(f"\nKL(pi* || R_01) = {endpoint_cost:.4f}")
-        print(f"  the integrand at t = 1/2 is {at_half:.4f}, and its integral is {integral:.4f}")
-        print(f"  at lambda = 1/2, 1, 2, 4 the cost is "
-              f"{[round(speeds[k], 4) for k in (0.5, 1.0, 2.0, 4.0)]}")
-        print(f"  the limit KL(pi_1 || (0.6, 0.4)) is {limit:.4f}")
-
-    report = {
-        "cost of halving every rate": scaling_cost(0.5),
-        "cost of doubling every rate": scaling_cost(2.0),
-        "R_01 column sums": [round(float(z), 4) for z in R01.sum(dim=0)],
-        "f": [round(float(z), 4) for z in f],
-        "g": [round(float(z), 4) for z in g],
-        "pi*_11": float(plan[0, 0]),
-        "phi at t = 1/2": [round(float(z), 4) for z in phi_half],
-        "bridge rate out of state 1": float(tilted_half[0, 1]),
-        "bridge rate out of state 2": float(tilted_half[1, 0]),
-        "rate out of state 1 at t = 1": float(tilted_one[0, 1]),
-        "marginal at t = 1/2": [round(float(z), 4) for z in marginal_half],
-        "marginal at t = 1": [round(float(z), 4) for z in marginal_one],
-        "reverse rate, 2 back to 1": float(reverse[1, 0]),
-        "Euler error at the coarsest step": sampler[coarse][1],
-        "Euler error, step halved": sampler[2 * coarse][1],
-        "Euler error, step quartered": sampler[4 * coarse][1],
-        "Euler error at ten times the steps": sampler[10 * coarse][1],
-        "KL(pi* || R_01)": endpoint_cost,
-        "rate error integrand at t = 1/2": at_half,
-        "its integral over [0, 1]": integral,
-        "cost at lambda = 1/2": speeds[0.5],
-        "cost at lambda = 4": speeds[4.0],
-        "KL(pi_1 || the stationary law)": limit,
-    }
-    print_report("The bridge over a finite set, Chapter 7 Sections 7.21 to 7.23.3", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/displacement.py

@@ -1,237 +0,0 @@
-"""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, and 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  # Closed forms throughout, so the printed digits are the chapter's.
-
-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), whose standard deviation is 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, which 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 attain the bound
-#
-# Drawing the endpoint pair from the optimal coupling and moving along the segment between them
-# makes the velocity constant, so Cauchy-Schwarz is 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, whose velocity is the constant x_1 - x_0."""
-    return (1.0 - t) * x0 + t * x1  # Same shape as x0 and x1 broadcast against t.
-
-
-# %% 05. Displacement interpolation
-#
-# Reading off the law of the particles at an intermediate time gives 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. Solving the
-# trajectory for X_0 and substituting makes the velocity 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, which is 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()

lecture_7/dsb.py

@@ -1,204 +0,0 @@
-"""Diffusion Schrodinger Bridge: iterative proportional fitting run on path laws.
-
-Each half-step of proportional fitting is a relative-entropy projection onto the set of path laws
-with one marginal correct, and the chain rule solves it in closed form by keeping the conditional
-law of the previous iterate and resetting that one marginal. The conditional law is the current
-process run backward in time, so a half-step amounts to reversing the current process and
-restarting it from the requested law at the other end.
-
-Chapter 7, Section 7.19.1 of the notes. Runs the half-steps exactly on the two-state chain with
-transition matrix A, source law (0.5, 0.5) and requested terminal law (0.2, 0.8), and checks the
-diffusion time reversal in closed form on a Gaussian. Prints the backward matrices, the law after
-each half-step, the forward matrices of the first round and the error at the free end.
-
-    python lecture_7/dsb.py              # 4 rounds, the first one worked by hand
-    python lecture_7/dsb.py --steps 1    # the single round the chapter writes out
-"""
-
-# %% 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  # The chapter quotes six figures of the backward and forward matrices.
-
-A = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=DTYPE)  # The reference transition matrix.
-PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE)  # The source law.
-PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE)  # The requested terminal law.
-TILTED_0 = torch.tensor([[0.599160, 0.400840], [0.299266, 0.700734]], dtype=DTYPE)  # A^h_0.
-TILTED_1 = torch.tensor([[0.307297, 0.692703], [0.112490, 0.887510]], dtype=DTYPE)  # A^h_1.
-
-
-# %% 02. A path law as a joint table
-#
-# With two states and two steps a path law is a table over the triple (x_0, x_1, x_2), so every
-# projection below is an exact operation on eight numbers. Section 7.12.1.
-
-
-def reference_path_law(start: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor:
-    """R(x_{0:N}) as a tensor with one axis per time, from the starting law and one matrix per step."""
-    law = start.clone()  # [K] the law at time 0.
-    for _ in range(steps):
-        law = law.unsqueeze(-1) * matrix  # Multiply by the transition matrix on a new axis.
-    return law  # [K] * (N + 1).
-
-
-def marginal(law: torch.Tensor, axis: int) -> torch.Tensor:
-    """The law of the state at one time, summed over every other time."""
-    other = [d for d in range(law.dim()) if d != axis]
-    return law.sum(dim=other)  # [K].
-
-
-def one_step_conditional(law: torch.Tensor, step: int) -> torch.Tensor:
-    """P(x_{k+1} | x_k), the forward transition matrix the path law carries at that step."""
-    pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)])  # [K, K].
-    return pair / pair.sum(dim=1, keepdim=True)  # Normalize each row.
-
-
-def reverse_step(law: torch.Tensor, step: int) -> torch.Tensor:
-    """P(x_k | x_{k+1}), Bayes' rule on one step, which is the reversal a network has to learn."""
-    pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)])  # [K, K].
-    return (pair / pair.sum(dim=0, keepdim=True)).T  # [K, K] indexed by the later state first.
-
-
-# %% 03. The two half-steps
-#
-# The odd iterate keeps the conditional law of the previous iterate given the final state and
-# replaces the marginal at time 1, and the even iterate does the same at time 0. Section 7.19.1.
-
-
-def fix_terminal(law: torch.Tensor, end: torch.Tensor) -> torch.Tensor:
-    """P'(x_{0:N}) = pi_1(x_N) P(x_{0:N-1} | x_N), the projection onto the laws ending at pi_1."""
-    last = marginal(law, law.dim() - 1)  # [K] the terminal marginal of the current law.
-    return law * (end / last)  # Reweight every path by the ratio at its final state.
-
-
-def fix_initial(law: torch.Tensor, start: torch.Tensor) -> torch.Tensor:
-    """P''(x_{0:N}) = pi_0(x_0) P'(x_{1:N} | x_0), the projection onto the laws starting at pi_0."""
-    first = marginal(law, 0)  # [K] the initial marginal of the current law.
-    shape = [-1] + [1] * (law.dim() - 1)  # Broadcast the ratio along the first axis.
-    return law * (start / first).reshape(shape)
-
-
-def proportional_fitting(law: torch.Tensor, start: torch.Tensor, end: torch.Tensor,
-                         rounds: int) -> tuple[torch.Tensor, list[float], list[torch.Tensor]]:
-    """Alternate the two half-steps, returning the law, the terminal error and the law each round."""
-    errors: list[float] = []
-    iterates: list[torch.Tensor] = []
-    for _ in range(rounds):
-        law = fix_terminal(law, end)  # One half-step: the terminal marginal is now exact.
-        law = fix_initial(law, start)  # The other: the initial marginal is now exact.
-        errors.append(float((marginal(law, law.dim() - 1) - end).abs().sum()))
-        iterates.append(law.clone())
-    return law, errors, iterates
-
-
-# %% 04. The reversal that a network learns
-#
-# On a finite state space the reversal is Bayes' rule on a matrix. On R^d it is the time reversal of
-# a diffusion, whose drift is minus the forward drift plus epsilon times the score of the current
-# marginal, and the score of the current iterate has no closed form. Section 7.19.1.
-
-
-def reverse_drift(x: torch.Tensor, t: torch.Tensor, epsilon: float, variance_0: float
-                  ) -> torch.Tensor:
-    """beta(x, t) = -b + epsilon grad log p_t for a driftless diffusion started from N(0, v_0)."""
-    variance = variance_0 + epsilon * t  # The marginal of sqrt(epsilon) dB started from N(0, v_0).
-    return epsilon * (-x / variance)  # The score of a centered Gaussian is -x over its variance.
-
-
-def reversal_check(epsilon: float, variance_0: float, steps: int, paths: int,
-                   generator: torch.Generator) -> tuple[float, float]:
-    """Run the reversal backward from p_1 and compare its variance at time 0 with v_0."""
-    dt = 1.0 / steps
-    variance_1 = variance_0 + epsilon * 1.0  # The forward marginal at time 1.
-    x = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(variance_1)
-    for step in range(steps):
-        t = torch.tensor(1.0 - step * dt, dtype=DTYPE)  # Time runs down from 1 to 0.
-        noise = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(dt)
-        x = x + reverse_drift(x, t, epsilon, variance_0) * dt + math.sqrt(epsilon) * noise
-    return float(x.var()), variance_0
-
-
-# %% 05. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=4)  # --steps is the number of proportional fitting rounds.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    rounds = max(args.steps, 1)
-    generator = torch.Generator().manual_seed(args.seed)
-
-    reference = reference_path_law(PI_0, A, 2)  # [2, 2, 2] the law P^0 = R.
-    marginals = [marginal(reference, k) for k in range(3)]
-    B2 = reverse_step(reference, 1)  # P(x_1 | x_2), the reversal of the last step.
-    B1 = reverse_step(reference, 0)  # P(x_0 | x_1), the reversal of the earlier step.
-    if loud:
-        print("The reference started from pi_0, read at the three time points")
-        print(f"  {[[round(float(z), 3) for z in m] for m in marginals]}")
-        print(f"  B_2 = {[[round(float(z), 6) for z in row] for row in B2]}")
-        print(f"  B_1(1,1) = {float(B1[0, 0]):.6f}")
-
-    half = fix_terminal(reference, PI_1)  # The first half-step, which resets the terminal marginal.
-    backward_laws = [marginal(half, 2), marginal(half, 1), marginal(half, 0)]
-    drift_at_zero = float((backward_laws[2] - PI_0).abs().sum())
-    full = fix_initial(half, PI_0)  # The second half-step, which resets the initial marginal.
-    F0 = one_step_conditional(full, 0)
-    F1 = one_step_conditional(full, 1)
-    terminal = marginal(full, 2)
-    terminal_error = float((terminal - PI_1).abs().sum())
-    tilt_gap = max(float((F0 - TILTED_0).abs().max()), float((F1 - TILTED_1).abs().max()))
-    if loud:
-        print("\nThe backward chain started from the requested terminal law")
-        written = ["(" + ", ".join(f"{float(z):.6f}" for z in m) + ")" for m in backward_laws]
-        print(f"  {' -> '.join(written)}")
-        print(f"  total absolute error at step 0: {drift_at_zero:.4f}")
-        print(f"\nThe forward matrices of the second half-step")
-        print(f"  F_0 = {[[round(float(z), 6) for z in row] for row in F0]}")
-        print(f"  F_1 = {[[round(float(z), 6) for z in row] for row in F1]}")
-        print(f"  terminal law {[round(float(z), 6) for z in terminal]}, "
-              f"off by {terminal_error:.4f}")
-        print(f"  largest gap against the tilted matrices A^h_0 and A^h_1: {tilt_gap:.4f}")
-
-    _, errors, _ = proportional_fitting(reference, PI_0, PI_1, rounds)
-    ratios = [errors[k] / errors[k + 1] for k in range(len(errors) - 1) if errors[k + 1] > 0]
-    if loud:
-        print(f"\nterminal error by round {['%.3e' % e for e in errors]}")
-        print(f"  ratios {[round(r, 1) for r in ratios]}")
-
-    simulated, exact = reversal_check(1.0, 1.0, 200, 8192, generator)
-    if loud:
-        print(f"\nThe diffusion reversal on a Gaussian, run backward from p_1")
-        print(f"  variance at time 0: simulated {simulated:.3f}, exact {exact:.3f}")
-
-    report = {
-        "reference marginal at step 1": [round(float(z), 3) for z in marginals[1]],
-        "reference marginal at step 2": [round(float(z), 3) for z in marginals[2]],
-        "B_2(1,1)": float(B2[0, 0]),
-        "B_2(2,1)": float(B2[1, 0]),
-        "B_1(1,1)": float(B1[0, 0]),
-        "backward law after one step": [round(float(z), 6) for z in backward_laws[1]],
-        "backward law at step 0": [round(float(z), 6) for z in backward_laws[2]],
-        "error at the free end, half-step": drift_at_zero,
-        "F_0(1,1)": float(F0[0, 0]),
-        "F_1(1,2)": float(F1[0, 1]),
-        "terminal law after one round": [round(float(z), 6) for z in terminal],
-        "terminal error after one round": terminal_error,
-        "gap against the tilted matrices": tilt_gap,
-        "terminal error, last round": errors[-1],
-        "error ratio per round": ratios[0] if ratios else float("nan"),
-        "rounds run": rounds,
-        "reversal variance, simulated": simulated,
-        "reversal variance, exact": exact,
-    }
-    print_report("Diffusion Schrodinger Bridge, Chapter 7 Section 7.19.1", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/dsbm.py

@@ -1,207 +0,0 @@
-"""Diffusion Schrodinger Bridge Matching: the Markov projection and the reciprocal projection.
-
-A mixture of bridges draws a pair of ends from a coupling and runs the reference bridge between
-them, so its drift depends on the endpoint and the law is in general not Markov. Averaging that
-drift over the endpoints of all bridges through a state removes the extra information and leaves a
-diffusion with the same marginals, and replacing the bridges of a Markov law by reference bridges
-keeps its endpoint coupling. Alternating the two projections holds both marginals fixed throughout
-and improves only the pairing.
-
-Chapter 7, Section 7.19.2 of the notes. Runs the two projections exactly on the two-state chain,
-started from the independent coupling of (0.5, 0.5) and (0.2, 0.8). Prints the averaged drift at one
-point, the two projected matrices of the first round, and the gap against the bridge coupling.
-
-    python lecture_7/dsbm.py              # 9 rounds, which reach 1e-9 on this problem
-    python lecture_7/dsbm.py --steps 1    # the single round the chapter writes out
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import torch
-
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-DTYPE = torch.float64  # The chapter quotes six figures of the two projected matrices.
-
-A = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=DTYPE)  # The reference transition matrix.
-PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE)  # The source law.
-PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE)  # The requested terminal law.
-
-
-# %% 02. The averaged bridge drift
-#
-# Two bridges passing through the same state at the same time carry different drifts, and the
-# Markov drift at that state is their average weighted by how often each one passes through.
-# A least-squares fit of a single number against the two targets lands on the same value.
-# Section 7.19.2.
-
-
-def bridge_drift(x1: float, x: float, t: float) -> float:
-    """(x_1 - x) / (1 - t), the drift of the Brownian bridge ending at x_1."""
-    return (x1 - x) / (1.0 - t)
-
-
-def markov_drift(endpoints: torch.Tensor, weights: torch.Tensor, x: float, t: float) -> float:
-    """v*(x, t) = E[(X_1 - x) / (1 - t) | X_t = x], the average over the bridges through x."""
-    drifts = (endpoints - x) / (1.0 - t)  # [n] one conditional drift per endpoint.
-    return float((weights * drifts).sum())  # The posterior average of the conditional drifts.
-
-
-def least_squares_derivative(a: float, endpoints: torch.Tensor, weights: torch.Tensor,
-                             x: float, t: float) -> float:
-    """The derivative in a of the average squared error of a single prediction against the drifts."""
-    drifts = (endpoints - x) / (1.0 - t)  # [n].
-    return float(2.0 * (weights * (a - drifts)).sum())  # Vanishes at the weighted average.
-
-
-# %% 03. The reciprocal projection
-#
-# A member of the reciprocal class is an endpoint coupling filled in with reference bridges. The
-# projection of a Markov law keeps that law's endpoint coupling and redraws the interior from the
-# reference. Section 7.19.2.
-
-
-def reference_bridge_table(matrix: torch.Tensor) -> torch.Tensor:
-    """R(x_1 | x_0, x_2) = A(x_0, x_1) A(x_1, x_2) / A^2(x_0, x_2) for every triple. [K, K, K]."""
-    two_step = matrix @ matrix  # [K, K] the reference over both steps.
-    joint = matrix[:, :, None] * matrix[None, :, :]  # [x_0, x_1, x_2] the unnormalized weight.
-    return joint / two_step[:, None, :]  # Normalize over the middle state.
-
-
-def reciprocal_projection(coupling: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
-    """proj_R(M) = M_{0,1} R_{.|0,1}, the coupling with the reference bridges put back in."""
-    return coupling[:, None, :] * reference_bridge_table(matrix)  # [K, K, K] a full path law.
-
-
-# %% 04. The Markov projection
-#
-# Retaining the one-step conditionals of a path law and multiplying them into a chain gives the
-# Markov law with the same marginals, and the averaging in each conditional is the discrete form of
-# the drift average. Section 7.19.2.
-
-
-def one_step_conditional(law: torch.Tensor, step: int) -> torch.Tensor:
-    """P(x_{k+1} | x_k) read off a path law, which is what the network fits at that transition."""
-    pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)])  # [K, K].
-    return pair / pair.sum(dim=1, keepdim=True)
-
-
-def markov_projection(law: torch.Tensor) -> tuple[list[torch.Tensor], torch.Tensor]:
-    """The chain of one-step conditionals, and the endpoint coupling it produces from pi_0."""
-    conditionals = [one_step_conditional(law, k) for k in range(law.dim() - 1)]
-    joint = torch.diag(marginal(law, 0))  # [K, K] the initial law on the diagonal.
-    for matrix in conditionals:
-        joint = joint @ matrix  # Multiply the conditionals into a chain.
-    return conditionals, joint
-
-
-def marginal(law: torch.Tensor, axis: int) -> torch.Tensor:
-    """The law of the state at one time, summed over every other time."""
-    return law.sum(dim=[d for d in range(law.dim()) if d != axis])
-
-
-# %% 05. Alternating the two projections
-#
-# Under a Markov projection every marginal is unchanged and under a reciprocal projection the
-# endpoint coupling is unchanged, so both pi_0 and pi_1 stay exact once the starting law has them,
-# and the iterations improve only the coupling. Section 7.19.2.
-
-
-def imf_rounds(coupling: torch.Tensor, matrix: torch.Tensor, rounds: int,
-               target: torch.Tensor) -> tuple[torch.Tensor, list[float], list]:
-    """Run iterative Markovian fitting, returning the coupling, the gaps and the first round."""
-    gaps: list[float] = []
-    first: list = []
-    for index in range(rounds):
-        law = reciprocal_projection(coupling, matrix)  # Reference bridges between the current ends.
-        conditionals, coupling = markov_projection(law)  # The Markov law and its endpoint coupling.
-        gaps.append(float((coupling - target).abs().sum()))
-        if index == 0:
-            first = conditionals
-    return coupling, gaps, first
-
-
-def bridge_coupling(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor,
-                    rounds: int = 200) -> torch.Tensor:
-    """The exact bridge coupling, by scaling diag(pi_0) A^2 to the two marginals."""
-    R01 = torch.diag(start) @ torch.linalg.matrix_power(matrix, 2)
-    f = torch.ones_like(start)
-    for _ in range(rounds):
-        g = end / (R01.T @ f)
-        f = start / (R01 @ g)
-    return f[:, None] * R01 * g[None, :]
-
-
-# %% 06. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=9)  # --steps is the number of Markovian fitting rounds.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    rounds = max(args.steps, 1)
-
-    endpoints = torch.tensor([3.0, 0.0], dtype=DTYPE)  # Two bridges through x = 1 at t = 0.5.
-    weights = torch.tensor([0.5, 0.5], dtype=DTYPE)  # They pass with equal probability.
-    average = markov_drift(endpoints, weights, 1.0, 0.5)
-    if loud:
-        print("Two bridges through x = 1 at t = 0.5, ending at 3 and at 0")
-        print(f"  their drifts {bridge_drift(3.0, 1.0, 0.5):.0f} and "
-              f"{bridge_drift(0.0, 1.0, 0.5):.0f}, average {average:.0f}")
-        print(f"  derivative of the squared error at a = {average:.0f}: "
-              f"{least_squares_derivative(average, endpoints, weights, 1.0, 0.5):.0f}")
-        print(f"  derivative at a = 0: "
-              f"{least_squares_derivative(0.0, endpoints, weights, 1.0, 0.5):+.0f}")
-
-    independent = PI_0[:, None] * PI_1[None, :]  # [[0.1, 0.4], [0.1, 0.4]].
-    target = bridge_coupling(A, PI_0, PI_1)
-    bridges = reference_bridge_table(A)
-    starting_gap = float((independent - target).abs().sum())
-    coupling, gaps, first = imf_rounds(independent, A, rounds, target)
-    projected = torch.diag(PI_0) @ first[0] @ first[1]
-    ratios = [gaps[k] / gaps[k + 1] for k in range(len(gaps) - 1) if gaps[k + 1] > 0]
-    if loud:
-        print(f"\nStarting from the independent coupling "
-              f"{[[round(float(z), 1) for z in row] for row in independent]}")
-        print(f"  the middle state is 1 with probability {float(bridges[0, 0, 0]):.6f} "
-              f"from 1 to 1, and {float(bridges[0, 0, 1]):.6f} from 1 to 2")
-        print(f"  M_0 = {[[round(float(z), 6) for z in row] for row in first[0]]}")
-        print(f"  M_1 = {[[round(float(z), 6) for z in row] for row in first[1]]}")
-        print(f"  M_0(1,1) = 0.2 x {float(bridges[0, 0, 0]):.6f} + 0.8 x "
-              f"{float(bridges[0, 0, 1]):.6f} = {float(first[0][0, 0]):.6f}")
-        print(f"  diag(pi_0) M_0 M_1 = "
-              f"{[[round(float(z), 6) for z in row] for row in projected]}")
-        print(f"  row sums {[round(float(z), 3) for z in projected.sum(1)]}, "
-              f"column sums {[round(float(z), 3) for z in projected.sum(0)]}")
-        print(f"\ngap against pi*: {starting_gap:.4f} at the start, "
-              f"then {['%.2e' % g for g in gaps[:5]]}")
-        print(f"  ratios {[round(r, 2) for r in ratios[:4]]}")
-
-    reached = next((k + 1 for k, g in enumerate(gaps) if g < 1e-9), None)
-    report = {
-        "drift of the bridge ending at 3": bridge_drift(3.0, 1.0, 0.5),
-        "drift of the bridge ending at 0": bridge_drift(0.0, 1.0, 0.5),
-        "averaged Markov drift": average,
-        "least squares derivative at a = 1": least_squares_derivative(
-            average, endpoints, weights, 1.0, 0.5),
-        "R(mid = 1 | 1, 1)": float(bridges[0, 0, 0]),
-        "R(mid = 1 | 1, 2)": float(bridges[0, 0, 1]),
-        "M_0(1,1)": float(first[0][0, 0]),
-        "M_1(1,2)": float(first[1][0, 1]),
-        "endpoint coupling, entry 11": float(projected[0, 0]),
-        "endpoint coupling, entry 12": float(projected[0, 1]),
-        "row sums after the projection": [round(float(z), 6) for z in projected.sum(dim=1)],
-        "column sums after the projection": [round(float(z), 6) for z in projected.sum(dim=0)],
-        "gap, independent coupling": starting_gap,
-        "gap after one round": gaps[0],
-        "gap after the last round": gaps[-1],
-        "gap ratio per round": ratios[-1] if ratios else float("nan"),
-        "rounds to 1e-9": reached if reached else rounds,
-    }
-    print_report("Diffusion Schrodinger Bridge Matching, Chapter 7 Section 7.19.2", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/dynamic_bridge.py

@@ -1,257 +0,0 @@
-"""A diffusion reference: Girsanov, the control problem, Hopf-Cole, and the Brownian bridge.
-
-Changing the drift of a diffusion and keeping its noise is the only change under which the two path
-laws stay comparable, and over one short step the two increments differ only in their mean. Summing
-those steps gives Girsanov's formula, averaging it under the controlled law turns relative entropy
-on paths into an integral of a squared drift, and the Hopf-Cole change of variables turns the
-resulting Hamilton-Jacobi-Bellman equation into the linear backward equation of the reference.
-
-Chapter 7, Sections 7.16 to 7.18 of the notes. Runs on a driftless diffusion in one dimension.
-Prints the two computations of the relative entropy of a constant control, the residual of the
-backward equation and of the Hamilton-Jacobi-Bellman equation, and the Brownian bridge between two
-fixed points.
-
-    python lecture_7/dynamic_bridge.py              # 200 Euler-Maruyama steps
-    python lecture_7/dynamic_bridge.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  # Closed forms and a Monte Carlo check, both in double precision.
-
-EPSILON = 1.0  # The reference noise variance per unit time.
-START, END = 0.0, 2.0  # The two points the bridge is pinned at.
-PATHS = 8192  # Trajectories in the Monte Carlo checks.
-
-
-# %% 02. One short step
-#
-# Over a step of length dt both increments are Gaussian with the covariance epsilon dt, and they
-# differ only in the mean, by u dt. The logarithm of the ratio of the two densities is the inner
-# product of w, the increment beyond the reference drift, with u over epsilon, minus the squared
-# norm of u times dt over 2 epsilon. Section 7.16.1.
-
-
-def log_density_ratio_step(w: torch.Tensor, u: torch.Tensor, dt: float,
-                           epsilon: float = EPSILON) -> torch.Tensor:
-    """log p^u / p^R over one step, with w = Delta X - b dt the part beyond the reference drift."""
-    return (w * u).sum(dim=-1) / epsilon - (u ** 2).sum(dim=-1) * dt / (2.0 * epsilon)
-
-
-# %% 03. Summing over the steps, and the average under the controlled law
-#
-# Adding the per-step contributions and letting dt go to zero gives Girsanov's formula, a martingale
-# term built from the noise and a drift term built from the size of the control. Averaging under
-# P^u leaves the two energy terms with coefficients 1/epsilon and -1/(2 epsilon).
-# Sections 7.16.2 and 7.16.3.
-
-
-def simulate_controlled(u: float, steps: int, paths: int, epsilon: float = EPSILON,
-                        generator: torch.Generator | None = None) -> tuple[torch.Tensor, torch.Tensor]:
-    """Euler-Maruyama for dX = u dt + sqrt(epsilon) dB, returning the endpoints and the log ratio.
-
-    The log ratio is accumulated step by step from the same Brownian increments the path uses, which
-    is how Section 7.27.3 accumulates it during training.
-    """
-    dt = 1.0 / steps
-    x = torch.zeros(paths, 1, dtype=DTYPE)  # [P, 1] every path starts at the origin.
-    log_ratio = torch.zeros(paths, dtype=DTYPE)  # [P] the Girsanov integral along each path.
-    control = torch.full((paths, 1), u, dtype=DTYPE)  # [P, 1] a control constant in space and time.
-    for _ in range(steps):
-        noise = torch.randn(paths, 1, dtype=DTYPE, generator=generator) * math.sqrt(dt)  # dB.
-        w = control * dt + math.sqrt(epsilon) * noise  # Delta X with no reference drift.
-        log_ratio = log_ratio + log_density_ratio_step(w, control, dt, epsilon)
-        x = x + w  # The controlled state after this step.
-    return x.squeeze(1), log_ratio
-
-
-def control_energy_entropy(u: float, epsilon: float = EPSILON) -> float:
-    """KL(P^u || R) = (1 / 2 epsilon) E int_0^1 ||u||^2 dt, for a control constant in time."""
-    return u ** 2 / (2.0 * epsilon)  # The integral of a constant over the unit interval.
-
-
-def gaussian_entropy_gap(shift: float, variance: float) -> float:
-    """KL between two Gaussians with a common variance, the squared difference of the means over 2v."""
-    return shift ** 2 / (2.0 * variance)
-
-
-# %% 04. The value function and the optimal control
-#
-# Bellman's principle and Ito's formula give the Hamilton-Jacobi-Bellman equation, in which only two
-# terms contain the control. Completing the square in those two isolates u* = -grad V, and
-# substituting it back leaves a closed equation for V with a squared gradient in it.
-# Sections 7.17 to 7.17.2.
-
-
-def optimal_control(value, x: float, t: float) -> float:
-    """u*(x, t) = -grad V_t(x), the control pointing downhill on the cost-to-go."""
-    xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
-    v = value(xs, torch.tensor(t, dtype=DTYPE))
-    return float(-torch.autograd.grad(v, xs)[0])
-
-
-def hjb_residual(value, x: float, t: float, epsilon: float = EPSILON) -> float:
-    """d_t V + (epsilon / 2) grad^2 V - (1/2) ||grad V||^2, with no reference drift."""
-    xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
-    ts = torch.tensor(t, dtype=DTYPE, requires_grad=True)
-    v = value(xs, ts)
-    dv_dt = torch.autograd.grad(v, ts, create_graph=True)[0]
-    dv_dx = torch.autograd.grad(v, xs, create_graph=True)[0]
-    d2v_dx2 = torch.autograd.grad(dv_dx, xs, create_graph=True)[0]
-    return float((dv_dt + 0.5 * epsilon * d2v_dx2 - 0.5 * dv_dx ** 2).detach())
-
-
-# %% 05. Removing the squared gradient
-#
-# Writing the value function as minus epsilon times a logarithm cancels the squared gradient against
-# the extra term the Laplacian produces, and what is left is the Kolmogorov backward equation of the
-# reference, which is linear. In the new variable the control is epsilon times the gradient of the
-# logarithm of phi. Sections 7.17.3 and 7.17.4.
-
-
-def hopf_cole(potential, epsilon: float = EPSILON):
-    """V_t = -epsilon log phi_t, the change of variables that linearizes the equation."""
-    return lambda x, t: -epsilon * torch.log(potential(x, t))
-
-
-def backward_residual(potential, x: float, t: float, epsilon: float = EPSILON) -> float:
-    """d_t phi + (epsilon / 2) grad^2 phi, which vanishes for a reference transition density."""
-    xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
-    ts = torch.tensor(t, dtype=DTYPE, requires_grad=True)
-    phi = potential(xs, ts)
-    dphi_dt = torch.autograd.grad(phi, ts, create_graph=True)[0]
-    dphi_dx = torch.autograd.grad(phi, xs, create_graph=True)[0]
-    d2phi_dx2 = torch.autograd.grad(dphi_dx, xs, create_graph=True)[0]
-    return float((dphi_dt + 0.5 * epsilon * d2phi_dx2).detach())
-
-
-# %% 06. The Brownian bridge
-#
-# Conditioning on a single final point y makes the potential the reference density of arriving at y
-# from the current state, and epsilon times the gradient of its logarithm is the remaining
-# displacement over the remaining time. The noise level cancels from the drift and survives in the
-# variance. Section 7.18.
-
-
-def bridge_potential(z: torch.Tensor, t: torch.Tensor, y: float = END,
-                     epsilon: float = EPSILON) -> torch.Tensor:
-    """h(z, t) = R_{1|t}(y | z), the full Gaussian density with variance epsilon (1 - t)."""
-    remaining = 1.0 - t  # The time left, which is the variance divided by epsilon.
-    return torch.exp(-((y - z) ** 2) / (2.0 * epsilon * remaining)) / torch.sqrt(
-        2.0 * math.pi * epsilon * remaining)
-
-
-def bridge_drift(z: torch.Tensor, t: torch.Tensor, y: float = END) -> torch.Tensor:
-    """(y - z) / (1 - t), the velocity that reaches y on time if nothing further perturbs the state."""
-    return (y - z) / (1.0 - t)
-
-
-def bridge_law(t: float, x: float = START, y: float = END,
-               epsilon: float = EPSILON) -> tuple[float, float]:
-    """The mean (1 - t) x + t y and the variance epsilon t (1 - t) of the bridge at time t."""
-    return (1.0 - t) * x + t * y, epsilon * t * (1.0 - t)
-
-
-def simulate_bridge(steps: int, paths: int, epsilon: float = EPSILON,
-                    generator: torch.Generator | None = None) -> torch.Tensor:
-    """Euler-Maruyama for the bridge SDE, returning the state at every time. [steps + 1, P]."""
-    dt = 1.0 / steps
-    x = torch.full((paths,), START, dtype=DTYPE)
-    trajectory = [x.clone()]
-    for step in range(steps):
-        t = torch.tensor(step * dt, dtype=DTYPE)
-        noise = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(dt)
-        x = x + bridge_drift(x, t) * dt + math.sqrt(epsilon) * noise  # One Euler-Maruyama step.
-        trajectory.append(x.clone())
-    return torch.stack(trajectory)  # [steps + 1, P].
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=200)  # --steps is the number of Euler-Maruyama steps.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    steps = max(args.steps, 2)
-    generator = torch.Generator().manual_seed(args.seed)
-
-    endpoints, log_ratio = simulate_controlled(2.0, steps, PATHS, EPSILON, generator)
-    girsanov = float(log_ratio.mean())  # The average of the log density under P^u.
-    closed_form = control_energy_entropy(2.0)
-    endpoint_form = gaussian_entropy_gap(2.0, EPSILON)
-    if loud:
-        print(f"A constant control u = 2 at epsilon = {EPSILON:g}, over {steps} steps")
-        print(f"  (1 / 2 eps) int ||u||^2 dt = {closed_form:.3f}")
-        print(f"  KL(N(x + 2, 1) || N(x, 1)) = {endpoint_form:.3f}")
-        print(f"  averaged Girsanov log ratio over {PATHS} paths = {girsanov:.3f}")
-        print(f"  endpoint mean {float(endpoints.mean()):.3f}, "
-              f"variance {float(endpoints.var()):.3f}")
-
-    value = hopf_cole(bridge_potential)  # V_t = -epsilon log h, with h the pinned potential.
-    test_points = [(-0.5, 0.2), (0.5, 0.5), (1.5, 0.5), (2.5, 0.8)]
-    backward = max(abs(backward_residual(bridge_potential, x, t)) for x, t in test_points)
-    hjb = max(abs(hjb_residual(value, x, t)) for x, t in test_points)
-    control_gap = max(
-        abs(optimal_control(value, x, t)
-            - float(bridge_drift(torch.tensor(x, dtype=DTYPE), torch.tensor(t, dtype=DTYPE))))
-        for x, t in test_points)
-    if loud:
-        print(f"\nThe pinned potential h(z, t) and its two equations")
-        print(f"  largest backward residual d_t h + (eps/2) h'' = {backward:.2e}")
-        print(f"  largest Hamilton-Jacobi-Bellman residual = {hjb:.2e}")
-        print(f"  largest gap between -grad V and (y - z) / (1 - t) = {control_gap:.2e}")
-
-    mean_half, variance_half = bridge_law(0.5)
-    sample = simulate_bridge(steps, PATHS, EPSILON, generator)
-    half = sample[steps // 2]  # The state at t = 1/2 on the grid.
-    state = 0.6732  # The particular state Section 7.18 evaluates the drift at.
-    drift_here = float(bridge_drift(torch.tensor(state, dtype=DTYPE),
-                                    torch.tensor(0.25, dtype=DTYPE)))
-    mean_quarter, variance_quarter = bridge_law(0.25)
-    spreads = {e: math.sqrt(e * 0.25) for e in (0.25, 1.0, 4.0)}  # Standard deviations at t = 1/2.
-    if loud:
-        print(f"\nThe bridge from {START:g} to {END:g} at epsilon = {EPSILON:g}")
-        print(f"  at t = 0.25: mean {mean_quarter:.4f}, variance {variance_quarter:.4f}, "
-              f"standard deviation {math.sqrt(variance_quarter):.4f}")
-        print(f"  at X_0.25 = {state}: drift = ({END:g} - {state}) / 0.75 = {drift_here:.4f}, "
-              f"against the straight-line velocity {END - START:.0f}")
-        print(f"  at t = 1/2: mean {mean_half:.4f}, variance {variance_half:.4f}")
-        print(f"  simulated mean {float(half.mean()):.4f}, variance {float(half.var()):.4f}")
-        print(f"  standard deviation at t = 1/2 for eps = 0.25, 1, 4: "
-              f"{[round(s, 2) for s in spreads.values()]}")
-
-    report = {
-        "KL from the control energy": closed_form,
-        "KL from the two endpoint laws": endpoint_form,
-        "KL from the simulated paths": girsanov,
-        "simulated endpoint mean": float(endpoints.mean()),
-        "backward equation residual": backward,
-        "Hamilton-Jacobi-Bellman residual": hjb,
-        "-grad V against the bridge drift": control_gap,
-        "bridge mean at t = 0.25": mean_quarter,
-        "bridge variance at t = 0.25": variance_quarter,
-        "bridge drift at X = 0.6732": drift_here,
-        "straight-line velocity": float(END - START),
-        "bridge mean at t = 1/2": mean_half,
-        "bridge variance at t = 1/2": variance_half,
-        "simulated mean at t = 1/2": float(half.mean()),
-        "simulated variance at t = 1/2": float(half.var()),
-        "simulated endpoint, mean": float(sample[-1].mean()),
-        "spread at t = 1/2, eps = 0.25": spreads[0.25],
-        "spread at t = 1/2, eps = 1": spreads[1.0],
-        "spread at t = 1/2, eps = 4": spreads[4.0],
-    }
-    print_report("A diffusion reference, Chapter 7 Sections 7.16 to 7.18", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/optimal_transport.py

@@ -1,323 +0,0 @@
-"""Optimal transport, from a pairing of four points to the Gaussian closed form.
-
-Two distributions fix how often each point appears and leave the assignment between them free. This
-file prices every assignment with the squared cost, writes a pairing as a matrix, relaxes Monge's
-map to Kantorovich's coupling, certifies the optimum with prices, and ends with the two Gaussians
-whose optimal map is an affine function.
-
-Chapter 7, Sections 7.1 to 7.6 of the notes. Runs on the chapter's own points, with no training and
-no data. Prints the cost of each pairing, the plan and the prices of the two-by-two example, the
-sorted pairing on three points, and the transport of N(0,1) to N(3,4).
-
-    python lecture_7/optimal_transport.py              # every worked number of Sections 7.1 to 7.6
-    python lecture_7/optimal_transport.py --steps 20   # a coarse quantile grid, for a first look
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import itertools
-import math
-
-import torch
-
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-DTYPE = torch.float64  # The chapter quotes six figures, so every number here is double precision.
-
-
-# %% 02. The cost of moving one unit of mass
-#
-# The price of carrying one unit from x to y is the squared Euclidean distance, under which a move
-# twice as long costs four times as much. Section 7.1.1.
-
-
-def cost_matrix(x: torch.Tensor, y: torch.Tensor, p: float = 2.0) -> torch.Tensor:
-    """C_ij = |x_i - y_j|^p on the line. x: [n], y: [m] -> [n, m]."""
-    return (x[:, None] - y[None, :]).abs() ** p  # [n, m] one price per route.
-
-
-def plan_cost(C: torch.Tensor, pi: torch.Tensor) -> float:
-    """The Frobenius inner product <C, pi>, the mass on each route times the price of that route."""
-    return float((C * pi).sum())
-
-
-# %% 03. Four points and two pairings
-#
-# Sources at 0 and 2, targets at 1 and 3, each point with mass 1/2. Section 7.1.2 counts the two
-# arrows of a pairing without weighting them by mass, and Section 7.1.3 weights them.
-
-
-def two_pairings() -> dict:
-    """The ordered and crossed pairings of the four points, by arrow count and by mass."""
-    x = torch.tensor([0.0, 2.0], dtype=DTYPE)  # Source points.
-    y = torch.tensor([1.0, 3.0], dtype=DTYPE)  # Target points.
-    C = cost_matrix(x, y)  # [[1, 9], [1, 1]] the four route prices.
-    ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=DTYPE)  # 0 -> 1 and 2 -> 3.
-    crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=DTYPE)  # 0 -> 3 and 2 -> 1.
-    independent = torch.full((2, 2), 0.25, dtype=DTYPE)  # a_i b_j with a = b = (1/2, 1/2).
-    return {
-        "C": C,
-        "arrows_ordered": float(C[0, 0] + C[1, 1]),  # 1 + 1, the two arrow lengths squared.
-        "arrows_crossed": float(C[0, 1] + C[1, 0]),  # 9 + 1.
-        "mass_ordered": plan_cost(C, ordered),  # Each arrow weighted by the 1/2 it carries.
-        "mass_crossed": plan_cost(C, crossed),
-        "mass_independent": plan_cost(C, independent),  # 1/4 on every route.
-    }
-
-
-# %% 04. A pairing written as a matrix
-#
-# A coupling is an n by m matrix with row sums a and column sums b, so pi 1_m = a and pi^T 1_n = b.
-# The independent coupling a_i b_j always lies in that set. Section 7.1.3.
-
-
-def is_coupling(pi: torch.Tensor, a: torch.Tensor, b: torch.Tensor, tol: float = 1e-9) -> bool:
-    """True when the row sums equal a, the column sums equal b, and no entry is negative."""
-    return bool(
-        (pi >= -tol).all()
-        and (pi.sum(dim=1) - a).abs().max() < tol  # pi 1_m = a.
-        and (pi.sum(dim=0) - b).abs().max() < tol  # pi^T 1_n = b.
-    )
-
-
-# %% 05. Monge's formulation, and the source that has to be split
-#
-# A map sends each source point to one destination, so for n equal masses it is a permutation and
-# the minimization runs over n! of them. A point mass pushes forward to a point mass, which is why
-# no map carries delta_0 to a law with mass at two places. Sections 7.2 and 7.2.1.
-
-
-def monge_permutations(x: torch.Tensor, y: torch.Tensor) -> list[tuple[float, tuple[int, ...]]]:
-    """Every permutation of n equal point masses and its average cost, cheapest first."""
-    n = len(x)
-    C = cost_matrix(x, y)  # [n, n].
-    costs = [(float(sum(C[i, p[i]] for i in range(n)) / n), p)
-             for p in itertools.permutations(range(n))]  # Mass 1/n on each of the n routes.
-    return sorted(costs)
-
-
-def pushforward_of_a_point_mass(destination: float) -> dict[float, float]:
-    """T_# delta_0 is the point mass at T(0), whatever the map T does elsewhere."""
-    return {destination: 1.0}  # One destination, all of the mass, so nu is out of reach.
-
-
-NU_ATOMS = {1.0: 0.5, 2.0: 0.5}  # nu = (1/2) delta_1 + (1/2) delta_2 of Section 7.2.1.
-
-
-# %% 06. Kantorovich's relaxation, solved by hand on the two-by-two
-#
-# The feasible set of the chapter's example is the one-parameter family pi(r), and the cost along it
-# is 5.8 - 8r, decreasing, so the optimum sits at the largest feasible r. Sections 7.3 to 7.3.2.
-
-SOURCE_X = torch.tensor([0.0, 2.0], dtype=DTYPE)  # Source points of the running example.
-SOURCE_A = torch.tensor([0.6, 0.4], dtype=DTYPE)  # Their masses.
-TARGET_Y = torch.tensor([1.0, 3.0], dtype=DTYPE)  # Target points.
-TARGET_B = torch.tensor([0.3, 0.7], dtype=DTYPE)  # Their required masses.
-
-
-def feasible_plan(r: float) -> torch.Tensor:
-    """The coupling with mass r from 0 to 1, the one degree of freedom the margins leave."""
-    return torch.tensor([[r, 0.6 - r], [0.3 - r, 0.1 + r]], dtype=DTYPE)  # [2, 2].
-
-
-def feasible_interval() -> tuple[float, float]:
-    """Nonnegativity of the four entries of pi(r) bounds r, and the binding pair is 0 <= r <= 0.3."""
-    lower = max(0.0, -0.1)  # From r >= 0 and 0.1 + r >= 0.
-    upper = min(0.6, 0.3)  # From 0.6 - r >= 0 and 0.3 - r >= 0.
-    return lower, upper
-
-
-def cost_along_the_segment(r: float, C: torch.Tensor) -> float:
-    """<C, pi(r)> as a function of r alone, which collects into 5.8 - 8r."""
-    return plan_cost(C, feasible_plan(r))  # Linear in r with slope C11 - C12 - C21 + C22 = -8.
-
-
-# %% 07. Prices and the dual
-#
-# A price f_i at each source point and g_j at each target point is feasible when f_i + g_j <= C_ij on
-# every route. Complementary slackness fixes the prices from the loaded routes of an optimal plan,
-# one substitution at a time, after one price is pinned to remove the shift. Sections 7.4 to 7.4.4.
-
-
-def prices_from_loaded_routes(C: torch.Tensor, pi: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """Solve f_i + g_j = C_ij on the routes carrying mass, with f_1 = 0 fixing the shift."""
-    n, m = C.shape
-    f = torch.full((n,), float("nan"), dtype=DTYPE)
-    g = torch.full((m,), float("nan"), dtype=DTYPE)
-    f[0] = 0.0  # Adding a constant to f and subtracting it from g changes no sum f_i + g_j.
-    for _ in range(n + m):  # Each pass propagates along one more loaded route.
-        for i in range(n):
-            for j in range(m):
-                if pi[i, j] <= 0:
-                    continue  # An empty route imposes no equation.
-                if not torch.isnan(f[i]) and torch.isnan(g[j]):
-                    g[j] = C[i, j] - f[i]  # The route prices the target point.
-                elif torch.isnan(f[i]) and not torch.isnan(g[j]):
-                    f[i] = C[i, j] - g[j]  # The route prices the source point.
-    return f, g
-
-
-def slack(C: torch.Tensor, f: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
-    """The matrix C_ij - f_i - g_j, nonnegative exactly when the prices are feasible."""
-    return C - f[:, None] - g[None, :]  # [n, m].
-
-
-def dual_value(a: torch.Tensor, f: torch.Tensor, b: torch.Tensor, g: torch.Tensor) -> float:
-    """<a, f> + <b, g>, the value of a pair of prices."""
-    return float(a @ f + b @ g)
-
-
-# %% 08. Powers of a distance, and non-crossing on the line
-#
-# With the squared cost the difference between the ordered and crossed pairings is -2 times the
-# product of the two gaps, so an optimal plan on the line contains no crossing. The quantile form
-# then computes W_p by sorting. Sections 7.5 and 7.5.1.
-
-
-def crossing_difference(x1: float, x2: float, y1: float, y2: float) -> tuple[float, float]:
-    """The ordered cost minus the crossed cost, expanded and in the collected form -2 dx dy."""
-    expanded = (x1 - y1) ** 2 + (x2 - y2) ** 2 - (x1 - y2) ** 2 - (x2 - y1) ** 2
-    collected = -2.0 * (x2 - x1) * (y2 - y1)  # Every squared term cancels, leaving the cross terms.
-    return expanded, collected
-
-
-def quantile_levels(levels: int) -> torch.Tensor:
-    """Midpoints (i + 1/2) / levels of a uniform grid on [0, 1], the levels q the integral uses."""
-    return (torch.arange(levels, dtype=DTYPE) + 0.5) / levels  # [levels].
-
-
-def empirical_quantile(points: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
-    """F^{-1}(q) = inf{x : F(x) >= q} for n equally weighted points, a step function of q."""
-    ordered = points.sort().values  # [n] the sorted points, which are the steps of F^{-1}.
-    index = torch.clamp((q * len(points)).floor().long(), max=len(points) - 1)  # [levels].
-    return ordered[index]  # [levels] the quantile at each level.
-
-
-def wasserstein_on_the_line(x: torch.Tensor, y: torch.Tensor, p: float, levels: int) -> float:
-    """W_p^p as the integral over quantile levels q of |F_mu^{-1}(q) - F_nu^{-1}(q)|^p.
-
-    With n equally weighted points the quantile function is a step function whose values are the
-    sorted points, so the level-by-level pairing is the sorted pairing and one sort on each side
-    computes the whole integral.
-    """
-    q = quantile_levels(levels)  # [levels].
-    difference = empirical_quantile(x, q) - empirical_quantile(y, q)  # [levels].
-    return float((difference.abs() ** p).mean())  # The integral of a step function over [0, 1].
-
-
-# %% 09. Gaussians, and the case where the plan is a map
-#
-# For two Gaussians on the line the squared 2-Wasserstein distance is the squared difference of the
-# means plus the squared difference of the standard deviations, and the optimal map is the increasing
-# affine function matching quantiles. Sections 7.6 to 7.6.2.
-
-
-def gaussian_w2_squared(m0: float, s0: float, m1: float, s1: float) -> float:
-    """W_2^2 between N(m0, s0^2) and N(m1, s1^2)."""
-    return (m0 - m1) ** 2 + (s0 - s1) ** 2
-
-
-def brenier_map(x: torch.Tensor, m0: float, s0: float, m1: float, s1: float) -> torch.Tensor:
-    """T(x) = m1 + (s1 / s0) (x - m0), the gradient of a convex function in one dimension."""
-    return m1 + (s1 / s0) * (x - m0)  # Same shape as x.
-
-
-# %% 10. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=3072)  # --steps is the number of quantile levels in Section 7.5.1.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-
-    pairings = two_pairings()
-    if loud:
-        print("Four points: sources at 0 and 2, targets at 1 and 3, mass 1/2 each")
-        print(f"  arrows   ordered {pairings['arrows_ordered']:.0f}   "
-              f"crossed {pairings['arrows_crossed']:.0f}")
-        print(f"  by mass  ordered {pairings['mass_ordered']:.0f}   "
-              f"crossed {pairings['mass_crossed']:.0f}   "
-              f"independent {pairings['mass_independent']:.0f}")
-
-    C = cost_matrix(SOURCE_X, TARGET_Y)  # [[1, 9], [1, 1]] for the running example.
-    low, high = feasible_interval()
-    grid = torch.linspace(low, high, 7, dtype=DTYPE)  # Seven points along the feasible segment.
-    best_r = float(grid[int(torch.tensor([cost_along_the_segment(float(r), C)
-                                          for r in grid]).argmin())])
-    pi_star = feasible_plan(high)  # The cost decreases in r, so the optimum is the right endpoint.
-    cost_star = plan_cost(C, pi_star)
-    if loud:
-        print(f"\nTwo-by-two: a = {SOURCE_A.tolist()}, b = {TARGET_B.tolist()}, C = {C.tolist()}")
-        print(f"  feasible r in [{low:.1f}, {high:.1f}],  cost(r) = "
-              f"{cost_along_the_segment(0.0, C):.1f} - "
-              f"{(cost_along_the_segment(0.0, C) - cost_along_the_segment(1.0, C)):.0f} r")
-        print(f"  pi* = {pi_star.tolist()}   <C, pi*> = {cost_star:.1f}")
-
-    f, g = prices_from_loaded_routes(C, pi_star)
-    S = slack(C, f, g)
-    dual = dual_value(SOURCE_A, f, TARGET_B, g)
-    gap = cost_star - dual  # Equals sum_ij pi_ij (C_ij - f_i - g_j), term by term nonnegative.
-    if loud:
-        print(f"\nPrices f = {f.tolist()}, g = {g.tolist()}")
-        print(f"  slack = {S.tolist()}   dual value {dual:.1f}   duality gap {gap:.1e}")
-        print(f"  slackness: mass on a route of positive slack = "
-              f"{float((pi_star * (S > 1e-12)).sum()):.1f}")
-
-    expanded, collected = crossing_difference(0.0, 2.0, 1.0, 3.0)
-    three_x = torch.tensor([4.0, 1.0, 6.0], dtype=DTYPE)  # Section 7.5.2.
-    three_y = torch.tensor([9.0, 2.0, 3.0], dtype=DTYPE)
-    permutations = monge_permutations(three_x.sort().values, three_y.sort().values)
-    quantile_w22 = wasserstein_on_the_line(three_x, three_y, 2.0, args.steps)
-    C3 = cost_matrix(three_x.sort().values, three_y.sort().values)
-    f3 = torch.tensor([0.0, 0.0, 8.0], dtype=DTYPE)  # Three equations in six unknowns, one solution.
-    g3 = torch.tensor([1.0, 1.0, 1.0], dtype=DTYPE)
-    dual3 = dual_value(torch.full((3,), 1 / 3, dtype=DTYPE), f3,
-                       torch.full((3,), 1 / 3, dtype=DTYPE), g3)
-    if loud:
-        print(f"\nNo crossing: expanded {expanded:+.0f} = collected {collected:+.0f}")
-        print(f"  three points, six permutations: {[round(c, 3) for c, _ in permutations]}")
-        print(f"  quantile integral on {args.steps} levels: {quantile_w22:.4f}")
-        print(f"  slack of f = (0, 0, 8) and g = (1, 1, 1): {slack(C3, f3, g3).tolist()}")
-        print(f"  dual value {dual3:.4f}")
-
-    source = torch.special.ndtri(quantile_levels(args.steps))  # N(0,1) read at the same levels.
-    mapped = brenier_map(source, 0.0, 1.0, 3.0, 2.0)  # T(x) = 3 + 2x.
-    displacement_square = float(((mapped - source) ** 2).mean())  # E[(3 + X)^2] = 9 + 1.
-    w22_gauss = gaussian_w2_squared(0.0, 1.0, 3.0, 2.0)
-    if loud:
-        probe = brenier_map(torch.tensor([-1.0, 0.0, 1.0], dtype=DTYPE), 0.0, 1.0, 3.0, 2.0)
-        print(f"\nN(0,1) to N(3,4): T(x) = {float(probe[1]):.0f} + "
-              f"{float(probe[2] - probe[1]):.0f}x, T(-1) = {float(probe[0]):.0f}, "
-              f"T(0) = {float(probe[1]):.0f}, T(1) = {float(probe[2]):.0f}")
-        print(f"  sampled mean square displacement {displacement_square:.3f}")
-
-    report = {
-        "ordered pairing, by arrows": pairings["arrows_ordered"],
-        "crossed pairing, by arrows": pairings["arrows_crossed"],
-        "ordered pairing, by mass": pairings["mass_ordered"],
-        "crossed pairing, by mass": pairings["mass_crossed"],
-        "independent coupling, by mass": pairings["mass_independent"],
-        "atoms of T_# delta_0, of nu": f"{len(pushforward_of_a_point_mass(1.0))}, {len(NU_ATOMS)}",
-        "permutations at n = 2, 3, 256": f"2, 6, a number with "
-                                         f"{len(str(math.factorial(256)))} digits",
-        "argmin of the cost over r": best_r,
-        "<C, pi*>": cost_star,
-        "dual value <a,f> + <b,g>": dual,
-        "duality gap": gap,
-        "pi* is a coupling": is_coupling(pi_star, SOURCE_A, TARGET_B),
-        "W_2 of the two-by-two": cost_star ** 0.5,
-        "sorted pairing on three points": permutations[0][0],
-        "worst pairing on three points": permutations[-1][0],
-        "W_2^2 from the quantile integral": quantile_w22,
-        "dual value on three points": dual3,
-        "W_2^2 of N(0,1) and N(3,4)": float(w22_gauss),
-        "W_2 of N(0,1) and N(3,4)": float(w22_gauss) ** 0.5,
-        "E[(T(X) - X)^2] on a sample": displacement_square,
-    }
-    print_report("Optimal transport, Chapter 7 Sections 7.1 to 7.6", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/ot_flow_matching.py

@@ -1,215 +0,0 @@
-"""Optimal transport inside flow matching: the same loss, trained on a different coupling.
-
-The flow matching loss of Chapter 2 is defined for any joint law whose two marginals are p_0 and
-p_1, and the independent draw of a source sample and a data sample is one such law. Replacing it
-with the optimal coupling of the batch changes the training pairs and leaves the loss, the network
-and the sampler as they are. Under the optimal coupling no two paths cross, so the marginal
-velocity at a point is the velocity of the single path through it and the trajectories come out
-straight.
-
-Chapter 7, Sections 7.8 and 7.8.1 of the notes. Trains two velocity networks on the four-Gaussian
-target of `cis6270.data`, one on the independent coupling and one on the exact coupling of each
-minibatch. Prints the coupling cost, the straightness of the learned paths and the one-step error.
-
-    python lecture_7/ot_flow_matching.py              # 300 steps per model, under a minute
-    python lecture_7/ot_flow_matching.py --steps 20   # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import torch
-
-from cis6270.data import gaussian_mixture_2d
-from cis6270.metrics import mode_report
-from cis6270.nets import VelocityMLP, count_parameters
-from cis6270.runner import (
-    common_parser, pick_device, print_report, save_run, set_seed, train,
-)
-
-
-# %% 02. The coupling of a batch
-#
-# A batch holds n source samples and m data samples with uniform masses 1/n and 1/m, and the
-# transport problem on that batch has nm unknowns against n + m equality constraints. With n = m the
-# vertex solution is a permutation, which the Hungarian algorithm returns in order n^3 operations.
-# Section 7.8.1.
-
-
-def batch_cost_matrix(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
-    """C_ij = ||x0_i - x1_j||^2 between every source sample and every data sample. [n, m]."""
-    return torch.cdist(x0, x1) ** 2  # [n, m] squared Euclidean cost.
-
-
-def optimal_pairing(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
-    """The permutation of the data samples that minimizes the total squared cost of the batch."""
-    from scipy.optimize import linear_sum_assignment  # The assignment problem of Section 7.2.
-
-    C = batch_cost_matrix(x0, x1).detach().cpu().numpy()  # [n, n].
-    _, columns = linear_sum_assignment(C)  # The cheapest permutation, in order n^3 operations.
-    return torch.as_tensor(columns, device=x0.device, dtype=torch.long)  # [n].
-
-
-def couple(x0: torch.Tensor, x1: torch.Tensor, kind: str) -> torch.Tensor:
-    """Reorder the data samples so that row i of the batch is the pair the coupling assigns to x0_i."""
-    if kind == "independent":
-        return x1  # Two independent draws, which is the coupling of Chapters 2 through 6.
-    return x1[optimal_pairing(x0, x1)]  # The optimal coupling of this batch.
-
-
-def vertex_count(n: int, m: int) -> dict[str, int]:
-    """The size of the batch problem: nm unknowns, n + m constraints, at most n + m - 1 loaded."""
-    return {"unknowns": n * m, "constraints": n + m, "loaded_routes": n + m - 1}
-
-
-# %% 03. The loss, with pi the only symbol that changes
-#
-# The network still predicts a velocity, the target is still the straight-line displacement, and the
-# time is still uniform on [0, 1]. Section 7.8.
-
-
-def flow_matching_loss(model: VelocityMLP, x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
-    """E_t ||v_theta((1-t) x_0 + t x_1) - (x_1 - x_0)||^2 over the pairs the coupling supplies."""
-    t = torch.rand(x0.shape[0], device=x0.device)  # [B] one time per pair.
-    xt = (1.0 - t[:, None]) * x0 + t[:, None] * x1  # [B, 2] the point on the straight path.
-    ut = x1 - x0  # [B, 2] the conditional velocity, constant along the path.
-    return ((model(xt, t) - ut) ** 2).sum(dim=1).mean()  # Sum coordinates, average pairs.
-
-
-def conditional_velocity_energy(x0: torch.Tensor, x1: torch.Tensor) -> float:
-    """The average squared conditional velocity, which Proposition 7.5 makes the kinetic action."""
-    return float(((x1 - x0) ** 2).sum(dim=1).mean())
-
-
-# %% 04. Sampling, and two measurements of the learned field
-#
-# A straight trajectory is integrated exactly by one Euler step, so the gap between a one-step
-# sample and a many-step sample measures how far the learned paths are from straight. The second
-# measurement compares the velocity along a trajectory with the net displacement of that
-# trajectory. Section 7.8.
-
-
-@torch.no_grad()
-def integrate(model: VelocityMLP, x0: torch.Tensor, steps: int) -> tuple[torch.Tensor, float]:
-    """Euler integration from t = 0 to t = 1, returning the samples and the path straightness.
-
-    The straightness is the average of ||v_t(x_t) - (x_1 - x_0)||^2 along the path, which is zero
-    exactly when the velocity never changes along a trajectory.
-    """
-    x = x0.clone()  # [B, 2].
-    dt = 1.0 / steps
-    velocities = []
-    for step in range(steps):
-        t = torch.full((x.shape[0],), step * dt, device=x.device)  # [B] shared time.
-        v = model(x, t)  # [B, 2] the learned velocity at the current state.
-        velocities.append(v)
-        x = x + dt * v  # One Euler step.
-    net = (x - x0)[None, :, :]  # [1, B, 2] the net displacement of each trajectory.
-    deviation = ((torch.stack(velocities) - net) ** 2).sum(dim=2).mean()  # Scalar.
-    return x, float(deviation)
-
-
-@torch.no_grad()
-def one_step_error(model: VelocityMLP, x0: torch.Tensor, steps: int) -> float:
-    """The distance between the one-step sample x_0 + v_0(x_0) and the many-step sample."""
-    fine, _ = integrate(model, x0, steps)  # [B, 2] the reference trajectory endpoint.
-    t0 = torch.zeros(x0.shape[0], device=x0.device)
-    coarse = x0 + model(x0, t0)  # [B, 2] a single Euler step of size 1.
-    return float((coarse - fine).norm(dim=1).mean())
-
-
-# %% 05. Training one model
-def fit(kind: str, args, device: torch.device, data: torch.Tensor) -> dict:
-    """Train a velocity network on the coupling named by ``kind`` and measure what it learned."""
-    set_seed(args.seed)  # The same initialization and the same batches for both couplings.
-    model = VelocityMLP(dim=2, width=args.width).to(device)
-    costs: list[float] = []
-
-    def loss_fn(step: int) -> torch.Tensor:
-        x1 = data[torch.randint(0, len(data), (args.batch_size,), device=device)]  # [B, 2].
-        x0 = torch.randn_like(x1)  # [B, 2] the source sample.
-        paired = couple(x0, x1, kind)  # [B, 2] reordered by the coupling.
-        costs.append(conditional_velocity_energy(x0, paired))
-        return flow_matching_loss(model, x0, paired)
-
-    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=True)
-    torch.manual_seed(args.seed + 1)  # The same starting noise for both samplers.
-    x0 = torch.randn(args.samples, 2, device=device)
-    samples, straightness = integrate(model, x0, args.sample_steps)
-    report = mode_report(samples.cpu())
-    return {
-        "kind": kind,
-        "loss": sum(losses[-50:]) / len(losses[-50:]),
-        "coupling_cost": sum(costs) / len(costs),
-        "straightness": straightness,
-        "one_step_error": one_step_error(model, x0, args.sample_steps),
-        "modes_found": report["modes_found"],
-        "nearest_distance": report["nearest_distance"],
-    }
-
-
-# %% 06. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=300, batch_size=128, width=64, samples=512, sample_steps=50)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    loud = not args.quiet
-    data = gaussian_mixture_2d(num_points=4096, seed=args.seed).to(device)  # [4096, 2].
-
-    # The two-point batch of Section 7.8.1, where both couplings can be written out.
-    x = torch.tensor([[0.0], [2.0]])  # Two source samples on the line.
-    y = torch.tensor([[1.0], [3.0]])  # Two data samples.
-    C = batch_cost_matrix(x, y)  # [[1, 9], [1, 1]].
-    independent_energy = float(C.mean())  # Weight 1/4 on each of the four pairs.
-    optimal_energy = float(C[torch.arange(2), optimal_pairing(x, y)].mean())  # Weight 1/2 each.
-    sizes = vertex_count(256, 256)  # The counting the chapter does, at n = m = 256.
-    if loud:
-        print("Two source samples at 0 and 2, two data samples at 1 and 3")
-        print(f"  independent coupling, average squared velocity {independent_energy:.0f}")
-        print(f"  optimal coupling, average squared velocity {optimal_energy:.0f}")
-        print(f"\nA batch of 256 and 256: {sizes['unknowns']:,} unknowns, "
-              f"{sizes['constraints']} constraints, at most {sizes['loaded_routes']} loaded routes")
-        print(f"  an exact solve costs about {256 ** 3:.3g} operations, repeated every batch")
-        print(f"  training batch {args.batch_size}, device={device}, parameters="
-              f"{count_parameters(VelocityMLP(dim=2, width=args.width)):,}")
-
-    results = {}
-    for kind in ("independent", "optimal"):
-        if loud:
-            print(f"\ntraining on the {kind} coupling")
-        results[kind] = fit(kind, args, device, data)
-        if loud:
-            r = results[kind]
-            print(f"  loss {r['loss']:.4f}   batch coupling cost {r['coupling_cost']:.3f}")
-            print(f"  straightness {r['straightness']:.4f}   "
-                  f"one-step error {r['one_step_error']:.4f}")
-
-    ratio = results["independent"]["coupling_cost"] / max(results["optimal"]["coupling_cost"], 1e-12)
-    report = {
-        "two-point batch, independent": independent_energy,
-        "two-point batch, optimal": optimal_energy,
-        "batch plan unknowns": sizes["unknowns"],
-        "batch plan constraints": sizes["constraints"],
-        "loaded routes at a vertex": sizes["loaded_routes"],
-        "coupling cost, independent": results["independent"]["coupling_cost"],
-        "coupling cost, optimal": results["optimal"]["coupling_cost"],
-        "cost ratio": ratio,
-        "training loss, independent": results["independent"]["loss"],
-        "training loss, optimal": results["optimal"]["loss"],
-        "straightness, independent": results["independent"]["straightness"],
-        "straightness, optimal": results["optimal"]["straightness"],
-        "one-step error, independent": results["independent"]["one_step_error"],
-        "one-step error, optimal": results["optimal"]["one_step_error"],
-        "modes found, independent": results["independent"]["modes_found"],
-        "modes found, optimal": results["optimal"]["modes_found"],
-        "nearest mode, independent": results["independent"]["nearest_distance"],
-        "nearest mode, optimal": results["optimal"]["nearest_distance"],
-    }
-    print_report("Optimal transport inside flow matching, Chapter 7 Section 7.8", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/sf2m.py

@@ -1,248 +0,0 @@
-"""Simulation-Free Score and Flow Matching: two closed-form targets at one bridge sample.
-
-With the entropic plan computed on a minibatch, the bridge is a mixture of Brownian bridges between
-pairs drawn from that plan, and a Brownian bridge at a time t is Gaussian. One normal draw gives a
-sample, and at that sample both the bridge drift and the score of the bridge marginal are available
-without simulating a path. Subtracting half the noise variance times the score turns the drift into
-the velocity of the probability flow, whose dependence on epsilon cancels.
-
-Chapter 7, Section 7.19.3 of the notes. Runs on the endpoint pair x_0 = 0, x_1 = 2 at epsilon = 1,
-and on the two-point batch of Section 7.8.1. Prints the two regression targets at t = 0.25 and at
-t = 0.75, the scaled score at every time on the grid, and the entropic plan of the batch.
-
-    python lecture_7/sf2m.py              # 101 times on the grid
-    python lecture_7/sf2m.py --steps 20   # a coarse time grid, for a first look
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import torch
-
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-DTYPE = torch.float64  # The chapter quotes four figures of both regression targets.
-
-X0, X1 = 0.0, 2.0  # The endpoint pair the chapter works through.
-EPSILON = 1.0  # The reference noise variance per unit time.
-DRAW = 0.4  # The normal draw z the chapter uses at both times.
-
-
-# %% 02. One sample from the bridge
-#
-# The Brownian bridge at time t is Gaussian with mean the straight interpolation and variance
-# epsilon t (1 - t), so a sample costs one normal draw and no intermediate step.
-# Section 7.19.3.
-
-
-def bridge_mean(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-    """m_t = (1 - t) x_0 + t x_1, the mean of the bridge between the two ends."""
-    return (1.0 - t) * x0 + t * x1
-
-
-def bridge_std(t: torch.Tensor, epsilon: float = EPSILON) -> torch.Tensor:
-    """sigma_t = sqrt(epsilon t (1 - t)), which vanishes at both ends and peaks at t = 1/2."""
-    return torch.sqrt(epsilon * t * (1.0 - t))
-
-
-def bridge_sample(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor, z: torch.Tensor,
-                  epsilon: float = EPSILON) -> torch.Tensor:
-    """x = m_t + sigma_t z, one draw from the bridge at time t."""
-    return bridge_mean(x0, x1, t) + bridge_std(t, epsilon) * z
-
-
-# %% 03. The two conditional quantities
-#
-# Both are conditional on the drawn pair, and neither depends on a quantity a network has to
-# estimate first. Section 7.19.3.
-
-
-def conditional_drift(x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-    """u_t(x | x_0, x_1) = (x_1 - x) / (1 - t), the remaining displacement over the remaining time."""
-    return (x1 - x) / (1.0 - t)
-
-
-def conditional_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
-                      epsilon: float = EPSILON) -> torch.Tensor:
-    """grad log p_t(x | x_0, x_1) = (m_t - x) / (epsilon t (1 - t)), displacement over variance."""
-    return (bridge_mean(x0, x1, t) - x) / (epsilon * t * (1.0 - t))
-
-
-# %% 04. From the bridge drift to a probability flow velocity
-#
-# Subtracting half the noise variance times the score converts a diffusion drift into the velocity
-# of the probability flow, which has the same marginals with no noise at all. The second form shows
-# the straight-line velocity plus a coefficient times the displacement of the sample from the mean,
-# with the dependence on epsilon cancelled. Section 7.19.3.
-
-
-def flow_coefficient(t: torch.Tensor) -> torch.Tensor:
-    """(1 - 2t) / (2 t (1 - t)), positive before t = 1/2 and negative after."""
-    return (1.0 - 2.0 * t) / (2.0 * t * (1.0 - t))
-
-
-def flow_velocity(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
-                  epsilon: float = EPSILON) -> torch.Tensor:
-    """u_t - (epsilon / 2) grad log p_t, the velocity the first network regresses onto."""
-    return (conditional_drift(x1, x, t)
-            - 0.5 * epsilon * conditional_score(x0, x1, x, t, epsilon))
-
-
-def flow_velocity_collected(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor,
-                            t: torch.Tensor) -> torch.Tensor:
-    """(x_1 - x_0) + (1 - 2t) / (2 t (1 - t)) (x - m_t), the same velocity with delta collected."""
-    return (x1 - x0) + flow_coefficient(t) * (x - bridge_mean(x0, x1, t))
-
-
-def scaled_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
-                 epsilon: float = EPSILON) -> torch.Tensor:
-    """sigma_t grad log p_t = (m_t - x) / sigma_t = -z, which has unit scale at every time."""
-    return bridge_std(t, epsilon) * conditional_score(x0, x1, x, t, epsilon)
-
-
-# %% 05. The loss and the sampler
-#
-# With lambda(t) = sigma_t the score target is -z, so no part of the schedule dominates the
-# gradient. At sampling time a noise level g(t) is fixed and the two networks are integrated
-# together, with g^2 = epsilon recovering the reference noise and g = 0 the probability flow.
-# Section 7.19.3.
-
-
-def sf2m_loss(velocity: torch.Tensor, score: torch.Tensor, velocity_target: torch.Tensor,
-              score_target: torch.Tensor, t: torch.Tensor, epsilon: float = EPSILON
-              ) -> torch.Tensor:
-    """||v_theta - u^o_t||^2 + lambda(t)^2 ||s_theta - grad log p_t||^2, with lambda(t) = sigma_t."""
-    weight = bridge_std(t, epsilon) ** 2  # lambda(t)^2, which turns the score target into -z.
-    return ((velocity - velocity_target) ** 2 + weight * (score - score_target) ** 2).mean()
-
-
-def sampler_drift(velocity: torch.Tensor, score: torch.Tensor, noise_level: float
-                  ) -> torch.Tensor:
-    """v_theta + (1/2) g(t)^2 s_theta, the drift of the family of samplers between the two ends."""
-    return velocity + 0.5 * noise_level ** 2 * score
-
-
-# %% 06. The role of the coupling
-#
-# Pairs from the entropic plan at strength 2 epsilon make the learned equation at g^2 = epsilon the
-# Schrodinger bridge, and pairs from any other coupling make it the Markov projection of that
-# mixture of bridges. The plan is computed on a minibatch. Section 7.19.3.
-
-
-def entropic_plan(x: torch.Tensor, y: torch.Tensor, a: torch.Tensor, b: torch.Tensor,
-                  epsilon: float = EPSILON, sweeps: int = 200) -> torch.Tensor:
-    """The entropic plan for the squared cost at strength 2 epsilon, by Sinkhorn on exp(-C / 2 eps)."""
-    C = (x[:, None] - y[None, :]) ** 2  # [n, m] the squared cost of the batch.
-    K = torch.exp(-C / (2.0 * epsilon))  # The Gibbs kernel a Brownian reference produces.
-    u, v = torch.ones_like(a), torch.ones_like(b)
-    for _ in range(sweeps):
-        u = a / (K @ v)
-        v = b / (K.T @ u)
-    return u[:, None] * K * v[None, :]  # [n, m].
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=101)  # --steps is the number of times on the grid.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    x0 = torch.tensor(X0, dtype=DTYPE)
-    x1 = torch.tensor(X1, dtype=DTYPE)
-    z = torch.tensor(DRAW, dtype=DTYPE)
-
-    worked = {}
-    for value in (0.25, 0.75):
-        t = torch.tensor(value, dtype=DTYPE)
-        x = bridge_sample(x0, x1, t, z)
-        worked[value] = {
-            "mean": float(bridge_mean(x0, x1, t)),
-            "std": float(bridge_std(t)),
-            "sample": float(x),
-            "drift": float(conditional_drift(x1, x, t)),
-            "score": float(conditional_score(x0, x1, x, t)),
-            "flow": float(flow_velocity(x0, x1, x, t)),
-            "collected": float(flow_velocity_collected(x0, x1, x, t)),
-            "scaled": float(scaled_score(x0, x1, x, t)),
-            "coefficient": float(flow_coefficient(t)),
-        }
-        if loud:
-            w = worked[value]
-            print(f"t = {value}, z = {DRAW}")
-            print(f"  m_t = {w['mean']:.4f}   sigma_t = {w['std']:.4f}   x = {w['sample']:.4f}")
-            print(f"  u_t = {w['drift']:.4f}   grad log p_t = {w['score']:.4f}")
-            print(f"  u^o_t = {w['flow']:.4f}   collected form {w['collected']:.4f}   "
-                  f"coefficient {w['coefficient']:+.4f}")
-            print(f"  sigma_t grad log p_t = {w['scaled']:.4f}, which is -z")
-
-    grid = torch.linspace(0.0, 1.0, max(args.steps, 3), dtype=DTYPE)[1:-1]  # Interior times.
-    samples = bridge_sample(x0, x1, grid, z)  # [levels] one bridge sample per time.
-    scaled = scaled_score(x0, x1, samples, grid)
-    scaled_gap = float((scaled + z).abs().max())  # The target is -z at every time.
-    flow_targets = flow_velocity(x0, x1, samples, grid)
-    score_targets = conditional_score(x0, x1, samples, grid)
-    loss_exact = float(sf2m_loss(flow_targets, score_targets, flow_targets, score_targets, grid))
-    loss_zero = float(sf2m_loss(torch.zeros_like(flow_targets), torch.zeros_like(score_targets),
-                                flow_targets, score_targets, grid))  # An untrained pair of networks.
-    stratified = torch.special.ndtri(
-        (torch.arange(max(args.steps, 3), dtype=DTYPE) + 0.5) / max(args.steps, 3))
-    half = torch.tensor(0.5, dtype=DTYPE)
-    averaged = float(flow_velocity(x0, x1, bridge_sample(x0, x1, half, stratified), half).mean())
-    symmetry = worked[0.25]["flow"] + worked[0.75]["flow"] - 2.0 * (X1 - X0)
-    if loud:
-        print(f"\nOn {len(grid)} interior times")
-        print(f"  largest gap between sigma_t grad log p_t and -z: {scaled_gap:.2e}")
-        print(f"  u^o averaged over the draw at t = 1/2: {averaged:.4f}, "
-              f"against the straight-line velocity {X1 - X0:.0f}")
-        print(f"  the two flow targets about the straight-line velocity: {symmetry:.2e}")
-        print(f"  loss at the exact targets {loss_exact:.2e}, at a zero network {loss_zero:.4f}")
-
-    x = torch.tensor([0.0, 2.0], dtype=DTYPE)  # The batch of Section 7.8.1.
-    y = torch.tensor([1.0, 3.0], dtype=DTYPE)
-    uniform = torch.full((2,), 0.5, dtype=DTYPE)
-    C = (x[:, None] - y[None, :]) ** 2
-    plan = entropic_plan(x, y, uniform, uniform)
-    ordered_cost = float(C[0, 0] + C[1, 1])  # 1 + 1, the arrows of the exact pairing.
-    crossed_cost = float(C[0, 1] + C[1, 0])  # 9 + 1, the arrows of the crossed pairing.
-    if loud:
-        print(f"\nThe two-point batch at epsilon = {EPSILON:g}")
-        print(f"  exact transport pairs 0 with 1 and 2 with 3 at cost {ordered_cost:.0f}, "
-              f"against {crossed_cost:.0f} for the crossed pairing")
-        print(f"  entropic plan at strength {2 * EPSILON:g}: "
-              f"{[[round(float(z_), 6) for z_ in row] for row in plan]}")
-        print(f"  mass on the crossed pairing {float(plan[0, 1] + plan[1, 0]):.6f}")
-
-    report = {
-        "m_t at t = 0.25": worked[0.25]["mean"],
-        "sigma_t at t = 0.25": worked[0.25]["std"],
-        "bridge sample at t = 0.25": worked[0.25]["sample"],
-        "u_t at t = 0.25": worked[0.25]["drift"],
-        "grad log p_t at t = 0.25": worked[0.25]["score"],
-        "u^o_t at t = 0.25": worked[0.25]["flow"],
-        "coefficient at t = 0.25": worked[0.25]["coefficient"],
-        "sigma_t grad log p_t at t = 0.25": worked[0.25]["scaled"],
-        "bridge sample at t = 0.75": worked[0.75]["sample"],
-        "u_t at t = 0.75": worked[0.75]["drift"],
-        "u^o_t at t = 0.75": worked[0.75]["flow"],
-        "coefficient at t = 0.75": worked[0.75]["coefficient"],
-        "the two targets, sum minus 2 u": symmetry,
-        "largest gap of the scaled score": scaled_gap,
-        "u^o averaged over the draw": averaged,
-        "loss at the exact targets": loss_exact,
-        "loss at a zero network": loss_zero,
-        "sampler drift at g^2 = epsilon": float(sampler_drift(
-            torch.tensor(worked[0.25]["flow"], dtype=DTYPE),
-            torch.tensor(worked[0.25]["score"], dtype=DTYPE), EPSILON ** 0.5)),
-        "sampler drift at g = 0": worked[0.25]["flow"],
-        "ordered pairing of the batch": ordered_cost,
-        "crossed pairing of the batch": crossed_cost,
-        "entropic plan, ordered mass": float(plan[0, 0] + plan[1, 1]),
-        "entropic plan, crossed mass": float(plan[0, 1] + plan[1, 0]),
-    }
-    print_report("Simulation-Free Score and Flow Matching, Chapter 7 Section 7.19.3", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/sinkhorn.py

@@ -1,234 +0,0 @@
-"""Entropic transport and Sinkhorn's algorithm on the chapter's two-by-two example.
-
-Subtracting epsilon times the entropy from the transport cost makes the objective strictly convex,
-and its minimizer is the Gibbs kernel exp(-C/epsilon) with one scaling per row and one per column.
-The two marginal constraints determine the scalings by alternate rescaling, which is one relative
-entropy projection per half-step.
-
-Chapter 7, Sections 7.9 and 7.10 of the notes. Runs on a = (0.6, 0.4), b = (0.3, 0.7) and the cost
-matrix with entries 1, 9, 1, 1. Prints the kernel, the first sweep by hand, the marginal error after
-each sweep, the converged plan and its multipliers, and the plan at four values of epsilon.
-
-    python lecture_7/sinkhorn.py                # 40 sweeps, every worked number of Section 7.10
-    python lecture_7/sinkhorn.py --steps 1      # the first sweep alone, computed by hand
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import torch
-
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-DTYPE = torch.float64  # The chapter quotes six figures of the kernel and of the converged plan.
-
-COST = torch.tensor([[1.0, 9.0], [1.0, 1.0]], dtype=DTYPE)  # C of Section 7.3.1.
-SOURCE_A = torch.tensor([0.6, 0.4], dtype=DTYPE)  # Row sums.
-TARGET_B = torch.tensor([0.3, 0.7], dtype=DTYPE)  # Column sums.
-EXACT_PLAN = torch.tensor([[0.3, 0.3], [0.0, 0.4]], dtype=DTYPE)  # pi* of Section 7.3.2.
-EPSILON = 2.0  # The strength the chapter runs by hand.
-
-
-# %% 02. Entropy of a plan, and the entropic objective
-#
-# Entropy is larger when the mass spreads over many routes and smaller when it concentrates, with
-# the convention 0 log 0 = 0. Subtracting epsilon H(pi) from the cost makes the objective strictly
-# convex, so the minimizer is unique and every entry of it is positive. Section 7.9.
-
-
-def entropy(pi: torch.Tensor) -> float:
-    """H(pi) = -sum_ij pi_ij log pi_ij, in nats, with an empty route contributing nothing."""
-    positive = pi[pi > 0]  # Drop the zero entries, which contribute 0 log 0 = 0.
-    return float(-(positive * positive.log()).sum())
-
-
-def entropic_objective(pi: torch.Tensor, C: torch.Tensor, epsilon: float) -> float:
-    """<C, pi> - epsilon H(pi), the quantity the entropic problem minimizes over Pi(a, b)."""
-    return float((C * pi).sum()) - epsilon * entropy(pi)
-
-
-# %% 03. Relative entropy from a positive matrix
-#
-# The definition keeps its meaning when xi is any strictly positive matrix. Taking xi to be the
-# independent coupling splits the logarithm into three pieces and gives -H(pi) + H(a) + H(b), so
-# the entropy term and a relative entropy from independence differ by a constant on Pi(a, b).
-# Section 7.9.1.
-
-
-def relative_entropy(pi: torch.Tensor, xi: torch.Tensor) -> float:
-    """KL(pi || xi) = sum_ij pi_ij log(pi_ij / xi_ij) for a strictly positive xi."""
-    mask = pi > 0  # Routes with no mass contribute nothing.
-    return float((pi[mask] * (pi[mask] / xi[mask]).log()).sum())
-
-
-def independent_coupling(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """a_i b_j, the coupling of largest entropy and the one Chapters 2 through 6 used."""
-    return a[:, None] * b[None, :]  # [n, m].
-
-
-# %% 04. The Lagrangian and the scaling form
-#
-# Setting the derivative of the Lagrangian in a single entry to zero gives
-# C_ij + epsilon (log pi_ij + 1) - f_i - g_j = 0, and solving for the logarithm splits the
-# minimizer into a row factor, a kernel entry and a column factor. Section 7.9.2.
-
-
-def gibbs_kernel(C: torch.Tensor, epsilon: float) -> torch.Tensor:
-    """K = exp(-C / epsilon), taken entrywise, with larger entries on the cheaper routes."""
-    return torch.exp(-C / epsilon)  # [n, m], every entry strictly positive.
-
-
-def scaling_form(u: torch.Tensor, K: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    """diag(u) K diag(v), the plan with row i scaled by u_i and column j scaled by v_j."""
-    return u[:, None] * K * v[None, :]  # [n, m].
-
-
-def multipliers_from_scalings(u: torch.Tensor, epsilon: float) -> torch.Tensor:
-    """Invert u_i = exp(f_i / epsilon - 1/2) for f_i = epsilon (log u_i + 1/2)."""
-    return epsilon * (u.log() + 0.5)  # [n] the additive price behind a multiplicative scaling.
-
-
-# %% 05. The two updates
-#
-# The row sums of diag(u) K diag(v) are u * (K v) and the column sums are v * (K^T u), and solving
-# each equation for its own scaling with the other held fixed gives the two updates. One sweep is
-# the pair of them, at two matrix-vector products. Section 7.10.
-
-
-def sinkhorn_sweeps(
-    K: torch.Tensor, a: torch.Tensor, b: torch.Tensor, sweeps: int,
-) -> tuple[torch.Tensor, torch.Tensor, list[float]]:
-    """Alternate u <- a / (K v) and v <- b / (K^T u), returning the scalings and the sweep errors."""
-    u = torch.ones_like(a)
-    v = torch.ones_like(b)
-    errors: list[float] = []
-    for _ in range(sweeps):
-        u = a / (K @ v)  # Row sums become exactly a.
-        v = b / (K.T @ u)  # Column sums become exactly b, and the row sums move away again.
-        errors.append(marginal_error(scaling_form(u, K, v), a, b))
-    return u, v, errors
-
-
-def marginal_error(pi: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> float:
-    """Total absolute error of the row sums against a plus that of the column sums against b."""
-    return float((pi.sum(dim=1) - a).abs().sum() + (pi.sum(dim=0) - b).abs().sum())
-
-
-def sweeps_to_tolerance(K: torch.Tensor, a: torch.Tensor, b: torch.Tensor,
-                        tolerance: float = 1e-9, limit: int = 2000) -> int:
-    """The number of sweeps that brings the total marginal error below the tolerance."""
-    _, _, errors = sinkhorn_sweeps(K, a, b, limit)
-    for index, error in enumerate(errors):
-        if error < tolerance:
-            return index + 1
-    return limit
-
-
-# %% 06. Each update as a relative entropy projection
-#
-# Minimizing KL(pi || xi) over the matrices with row sums a, with no condition on the columns, gives
-# pi_ij = xi_ij exp(lambda_i - 1), which is xi with each row rescaled, and the row sums fix the
-# factor at a_i / sum_j xi_ij. The u update performs exactly that rescaling. Section 7.10.1.
-
-
-def project_rows(xi: torch.Tensor, a: torch.Tensor) -> torch.Tensor:
-    """The matrix with row sums a closest to xi in relative entropy."""
-    return xi * (a / xi.sum(dim=1))[:, None]  # One factor per row.
-
-
-def project_columns(xi: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """The matrix with column sums b closest to xi in relative entropy."""
-    return xi * (b / xi.sum(dim=0))[None, :]  # One factor per column.
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--epsilon", type=float, default=EPSILON, help="regularization strength")
-    parser.set_defaults(steps=40)  # --steps is the number of Sinkhorn sweeps.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    epsilon = args.epsilon
-
-    K = gibbs_kernel(COST, epsilon)
-    independent = independent_coupling(SOURCE_A, TARGET_B)
-    if loud:
-        print(f"H(pi*) = {entropy(EXACT_PLAN):.3f}   H(a x b) = {entropy(independent):.3f}   "
-              f"H(a) + H(b) = {entropy(SOURCE_A) + entropy(TARGET_B):.3f}")
-        print(f"KL(pi* || a x b) = {relative_entropy(EXACT_PLAN, independent):.3f} nats")
-        print(f"\nK = exp(-C/{epsilon:g}) = {[[round(float(k), 6) for k in row] for row in K]}")
-        print(f"  kernel ratio between a cheap route and the expensive one "
-              f"{float(K[0, 0] / K[0, 1]):.1f}")
-
-    # The first sweep, written out the way Section 7.10.2 writes it.
-    v = torch.ones_like(TARGET_B)
-    Kv = K @ v  # Row sums of K, since v = (1, 1).
-    u = SOURCE_A / Kv
-    Ktu = K.T @ u
-    v = TARGET_B / Ktu
-    first = scaling_form(u, K, v)
-    if loud:
-        print(f"\nsweep 1   K v = {[round(float(z), 6) for z in Kv]}"
-              f"   u = {[round(float(z), 6) for z in u]}")
-        print(f"          K^T u = {[round(float(z), 6) for z in Ktu]}"
-              f"   v = {[round(float(z), 5) for z in v]}")
-        print(f"  diag(u) K diag(v) = {[[round(float(z), 4) for z in row] for row in first]}")
-        print(f"  column sums {[round(float(z), 4) for z in first.sum(0)]}  "
-              f"row sums {[round(float(z), 4) for z in first.sum(1)]}")
-        after_u = project_rows(first, SOURCE_A)  # The next update of u, as a row projection.
-        print(f"  after the next u update: rows {[round(float(z), 3) for z in after_u.sum(1)]}  "
-              f"columns {[round(float(z), 3) for z in after_u.sum(0)]}")
-
-    u, v, errors = sinkhorn_sweeps(K, SOURCE_A, TARGET_B, max(args.steps, 1))
-    plan = scaling_form(u, K, v)
-    f = multipliers_from_scalings(u, epsilon)
-    g = multipliers_from_scalings(v, epsilon)
-    sums = f[:, None] + g[None, :]  # The entropic counterparts of the dual prices.
-    if loud:
-        print(f"\nmarginal error by sweep {[round(e, 3) for e in errors[:5]]}")
-        print(f"  successive ratios "
-              f"{[round(errors[i + 1] / errors[i], 2) for i in range(min(4, len(errors) - 1))]}")
-        print(f"  pi^eps = {[[round(float(z), 6) for z in row] for row in plan]}")
-        print(f"  u = {[round(float(z), 4) for z in u]}   v = {[round(float(z), 4) for z in v]}")
-        print(f"  f = {[round(float(z), 3) for z in f]}   g = {[round(float(z), 3) for z in g]}")
-        print(f"  f_i + g_j = {[round(float(z), 3) for z in sums.flatten()]}  "
-              f"against C = {[round(float(z), 0) for z in COST.flatten()]}")
-        print(f"  the route (2,2) carries {float(plan[1, 1]):.3f} and overshoots its cost by "
-              f"{float(sums[1, 1] - COST[1, 1]):.3f}")
-
-        print("\neps     <C, pi^eps>    pi^eps_21        sweeps to 1e-9")
-        for value in (10.0, 2.0, 0.5, 0.1):
-            kernel = gibbs_kernel(COST, value)
-            su, sv, _ = sinkhorn_sweeps(kernel, SOURCE_A, TARGET_B, 400)
-            p = scaling_form(su, kernel, sv)
-            print(f"{value:<7g} {float((COST * p).sum()):<14.7f} {float(p[1, 0]):<16.3e} "
-                  f"{sweeps_to_tolerance(kernel, SOURCE_A, TARGET_B)}")
-
-    report = {
-        "H(pi*)": entropy(EXACT_PLAN),
-        "H(a x b) = H(a) + H(b)": entropy(independent),
-        "KL(pi* || a x b)": relative_entropy(EXACT_PLAN, independent),
-        "K_11": float(K[0, 0]),
-        "K_12": float(K[0, 1]),
-        "sweep 1 marginal error": errors[0],
-        "sweep 5 marginal error": errors[min(4, len(errors) - 1)],
-        "sweeps to 1e-9": sweeps_to_tolerance(K, SOURCE_A, TARGET_B),
-        "pi^eps_11": float(plan[0, 0]),
-        "pi^eps_21": float(plan[1, 0]),
-        "<C, pi^eps>": float((COST * plan).sum()),
-        "<C, pi*>": float((COST * EXACT_PLAN).sum()),
-        "H(pi^eps)": entropy(plan),
-        "objective at pi^eps": entropic_objective(plan, COST, epsilon),
-        "objective at pi*": entropic_objective(EXACT_PLAN, COST, epsilon),
-        "objective at a x b": entropic_objective(independent, COST, epsilon),
-        "KL(pi^eps || a x b)": relative_entropy(plan, independent),
-        "f": [round(float(z), 3) for z in f],
-        "g": [round(float(z), 3) for z in g],
-    }
-    print_report("Entropic transport and Sinkhorn, Chapter 7 Sections 7.9 and 7.10", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()

lecture_7/static_bridge.py

@@ -1,295 +0,0 @@
-"""The static Schrodinger bridge: a path law, the chain rule, and the scaling that solves it.
-
-A path law over a finite time grid is an ordinary probability vector over paths, and its relative
-entropy from a reference splits by the chain rule into a term for the two endpoints and a term for
-the interior. Only the first term is constrained, so the bridge keeps the reference bridges exactly
-and all of its freedom goes into the endpoint coupling, which iterative proportional fitting
-computes by alternate rescaling.
-
-Chapter 7, Sections 7.12 to 7.15.3 of the notes. Runs on the two-state chain with transition matrix
-A, source law (0.5, 0.5) and requested terminal law (0.2, 0.8). Prints the eight path probabilities,
-the chain rule on two tables, the fitted potentials, the tilted matrices and the bridge path law.
-
-    python lecture_7/static_bridge.py              # 4 rounds of proportional fitting
-    python lecture_7/static_bridge.py --steps 1    # one round, the one worked by hand
-"""
-
-# %% 01. Imports and configuration
-from __future__ import annotations
-
-import itertools
-
-import torch
-
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-DTYPE = torch.float64  # The chapter quotes six figures of the potentials and of the path law.
-
-A = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=DTYPE)  # The reference transition matrix.
-PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE)  # The source law.
-PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE)  # The requested terminal law after two steps.
-STEPS = 2  # N, the number of transitions, so a path is a triple and there are eight of them.
-
-
-# %% 02. The probability of a whole path
-#
-# For a Markov process the probability of a list of states is the starting mass times the product of
-# the transition probabilities along it. Section 7.12.1.
-
-
-def path_law(start: torch.Tensor, matrices: list[torch.Tensor]) -> dict[tuple[int, ...], float]:
-    """R(x_{0:N}) over every path of the chain, as a dictionary keyed by the list of states."""
-    size = len(start)
-    law = {}
-    for path in itertools.product(range(size), repeat=len(matrices) + 1):
-        weight = float(start[path[0]])  # The starting mass R_0(x_0).
-        for k, matrix in enumerate(matrices):
-            weight *= float(matrix[path[k], path[k + 1]])  # One transition probability per step.
-        law[path] = weight
-    return law  # K^(N+1) entries, summing to one.
-
-
-def relative_entropy(p: dict, q: dict) -> float:
-    """KL(P || R) = sum over paths of P log(P / R), the definition of Chapter 1 indexed by paths."""
-    return sum(v * torch.log(torch.tensor(v / q[k], dtype=DTYPE)).item()
-               for k, v in p.items() if v > 0)
-
-
-# %% 03. The chain rule for relative entropy
-#
-# Factoring each joint law into a marginal and a conditional and averaging the logarithm of the
-# ratio splits the relative entropy into a marginal term and an averaged conditional term.
-# Sections 7.13 and 7.13.1.
-
-
-def kl_vectors(p: torch.Tensor, q: torch.Tensor) -> float:
-    """Relative entropy between two probability vectors over the same states."""
-    mask = p > 0
-    return float((p[mask] * (p[mask] / q[mask]).log()).sum())
-
-
-def chain_rule_terms(pi: torch.Tensor, other: torch.Tensor) -> tuple[float, float, float]:
-    """The total, the marginal term and the averaged conditional term of KL(pi || other)."""
-    total = kl_vectors(pi.flatten(), other.flatten())  # The four cells, computed directly.
-    marginal = kl_vectors(pi.sum(dim=1), other.sum(dim=1))  # The disagreement about the row.
-    rows = [kl_vectors(pi[i] / pi[i].sum(), other[i] / other[i].sum()) for i in range(len(pi))]
-    conditional = float((pi.sum(dim=1) * torch.tensor(rows, dtype=DTYPE)).sum())
-    return total, marginal, conditional
-
-
-# %% 04. Splitting a path law at its two endpoints
-#
-# With the pair (x_0, x_N) in the place of the first coordinate and the interior in the place of the
-# second, the chain rule separates the endpoint coupling from the bridges. The second term is free
-# and nonnegative, so the bridge takes the reference bridges and the problem collapses onto the
-# coupling. Section 7.13.2.
-
-
-def endpoint_coupling(law: dict) -> torch.Tensor:
-    """The law of the pair (x_0, x_N) under a path law, summed over the interior states."""
-    coupling = torch.zeros(2, 2, dtype=DTYPE)
-    for path, weight in law.items():
-        coupling[path[0], path[-1]] += weight  # Collect every path with the same two ends.
-    return coupling
-
-
-def reference_bridge(x0: int, mid: int, xN: int, matrix: torch.Tensor,
-                     two_step: torch.Tensor) -> float:
-    """R(x_1 | x_0, x_2) = A(x_0, x_1) A(x_1, x_2) / A^2(x_0, x_2), the interior given the ends."""
-    return float(matrix[x0, mid] * matrix[mid, xN] / two_step[x0, xN])
-
-
-def fill_with_reference_bridges(coupling: torch.Tensor, matrix: torch.Tensor) -> dict:
-    """P(x_{0:2}) = pi(x_0, x_2) R(x_1 | x_0, x_2), a coupling filled in with reference bridges."""
-    two_step = matrix @ matrix
-    return {(x0, mid, xN): float(coupling[x0, xN]) * reference_bridge(x0, mid, xN, matrix, two_step)
-            for x0, mid, xN in itertools.product(range(2), repeat=3)}
-
-
-# %% 05. The static problem and the two potentials
-#
-# The bridge coupling minimizes KL(pi || R_01) over the couplings of pi_0 and pi_1, which is the
-# entropic problem with R_01 in place of the Gibbs kernel, so the minimizer is R_01 with one scaling
-# per row and one per column and the two marginals determine them by alternation.
-# Sections 7.14 and 7.14.2.
-
-
-def reference_endpoint_law(start: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor:
-    """R_01 = diag(pi_0) A^N, the reference probability of starting at x and ending at y."""
-    return torch.diag(start) @ torch.linalg.matrix_power(matrix, steps)  # [K, K].
-
-
-def proportional_fitting(
-    R01: torch.Tensor, start: torch.Tensor, end: torch.Tensor, rounds: int,
-) -> tuple[torch.Tensor, torch.Tensor, list[float]]:
-    """Alternate g <- pi_1 / (R_01^T f) and f <- pi_0 / (R_01 g), returning the terminal errors.
-
-    One round is a g update followed by an f update, started from f = (1, 1), so the row sums are
-    exact at the end of a round and the error is measured on the columns.
-    """
-    f = torch.ones_like(start)
-    g = torch.ones_like(end)
-    errors: list[float] = []
-    for _ in range(rounds):
-        g = end / (R01.T @ f)  # Column sums become exactly pi_1.
-        f = start / (R01 @ g)  # Row sums become exactly pi_0, and the columns move again.
-        plan = f[:, None] * R01 * g[None, :]
-        errors.append(float((plan.sum(dim=0) - end).abs().sum()))
-    return f, g, errors
-
-
-# %% 06. The squared cost from a Brownian reference
-#
-# For a Brownian motion of variance epsilon the transition density is a Gibbs kernel for the squared
-# distance at strength 2 epsilon, so minimizing the relative entropy from its endpoint law is the
-# entropic transport problem with the noise variance in the place of the regularization strength.
-# Section 7.14.1.
-
-
-def brownian_reference_plan(epsilon: float = 1.0) -> torch.Tensor:
-    """The bridge coupling of the four points of Section 7.3.1 under a Brownian reference."""
-    x = torch.tensor([0.0, 2.0], dtype=DTYPE)  # Source points.
-    y = torch.tensor([1.0, 3.0], dtype=DTYPE)  # Target points.
-    a = torch.tensor([0.6, 0.4], dtype=DTYPE)  # Their masses.
-    b = torch.tensor([0.3, 0.7], dtype=DTYPE)
-    endpoint_law = torch.exp(-(x[:, None] - y[None, :]) ** 2 / (2.0 * epsilon))  # exp(-C/2).
-    f, g, _ = proportional_fitting(endpoint_law, a, b, 200)
-    return f[:, None] * endpoint_law * g[None, :]  # [2, 2].
-
-
-# %% 07. Making the bridge a Markov process
-#
-# The backward potential averages g over the reference, one matrix-vector product per step, and
-# reweighting each transition by the ratio of the potential at the destination to the potential at
-# the current state produces a transition matrix whose rows sum to one by that same recursion.
-# Sections 7.15 and 7.15.1.
-
-
-def backward_potential(matrix: torch.Tensor, g: torch.Tensor, steps: int) -> list[torch.Tensor]:
-    """h_N = g and h_k = A h_{k+1}, returned as the list h_0, ..., h_N."""
-    potentials = [g]
-    for _ in range(steps):
-        potentials.insert(0, matrix @ potentials[0])  # One step back, into the front of the list.
-    return potentials
-
-
-def tilted_matrix(matrix: torch.Tensor, h_now: torch.Tensor, h_next: torch.Tensor) -> torch.Tensor:
-    """A^h_k(x, y) = A(x, y) h_{k+1}(y) / h_k(x), Doob's h-transform of one step."""
-    return matrix * h_next[None, :] / h_now[:, None]  # [K, K], each row summing to one.
-
-
-# %% 08. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=4)  # --steps is the number of proportional fitting rounds.
-    args = parser.parse_args()
-    set_seed(args.seed)
-    loud = not args.quiet
-    rounds = max(args.steps, 1)
-
-    reference = path_law(PI_0, [A] * STEPS)
-    if loud:
-        print("The eight paths of the two-state chain, with their reference probabilities")
-        for path, weight in reference.items():
-            print(f"  {','.join(str(s + 1) for s in path)}   R = {weight:.6f}")
-        print(f"  they sum to {sum(reference.values()):.6f}")
-
-    table = torch.tensor([[0.5, 0.1], [0.1, 0.3]], dtype=DTYPE)  # pi of Section 7.13.1.
-    uniform = torch.full((2, 2), 0.25, dtype=DTYPE)  # pi', the uniform law on the four cells.
-    total, marginal, conditional = chain_rule_terms(table, uniform)
-    if loud:
-        print(f"\nChain rule on two tables: total {total:.4f} = marginal {marginal:.4f} "
-              f"+ conditional {conditional:.4f}")
-
-    R01 = reference_endpoint_law(PI_0, A, STEPS)
-    reference_error = float((R01.sum(dim=0) - PI_1).abs().sum())
-    f, g, errors = proportional_fitting(R01, PI_0, PI_1, rounds)
-    plan = f[:, None] * R01 * g[None, :]
-    independent = PI_0[:, None] * PI_1[None, :]
-    if loud:
-        print(f"\nA^2 = {[[round(float(z), 2) for z in row] for row in A @ A]}")
-        print(f"R_01 = {[[round(float(z), 3) for z in row] for row in R01]}   "
-              f"column sums {[round(float(z), 3) for z in R01.sum(0)]}   error {reference_error:.2f}")
-        first_g = PI_1 / R01.sum(dim=0)  # The first update, with f = (1, 1).
-        print(f"  first g = {[round(float(z), 4) for z in first_g]}   "
-              f"R_01 g = {[round(float(z), 4) for z in R01 @ first_g]}")
-        print(f"  first f = {[round(float(z), 4) for z in PI_0 / (R01 @ first_g)]}")
-        print("  round   f                        g                        terminal error   ratio")
-        fr = torch.ones_like(PI_0)
-        for index in range(rounds):
-            gr = PI_1 / (R01.T @ fr)
-            fr = PI_0 / (R01 @ gr)
-            previous = reference_error if index == 0 else errors[index - 1]
-            ratio = f"{previous / errors[index]:.1f}" if errors[index] > 0 else "exact"
-            print(f"  {index + 1:<7d} ({float(fr[0]):.6f}, {float(fr[1]):.6f})   "
-                  f"({float(gr[0]):.6f}, {float(gr[1]):.6f})   {errors[index]:.3e}      {ratio}")
-        print(f"  pi* = {[[round(float(z), 6) for z in row] for row in plan]}")
-        print(f"  KL(pi* || R_01) = {kl_vectors(plan.flatten(), R01.flatten()):.4f}   "
-              f"KL(independent || R_01) = {kl_vectors(independent.flatten(), R01.flatten()):.4f}")
-
-    brownian = brownian_reference_plan()
-    if loud:
-        print("\nBrownian reference at epsilon = 1 on the four points of Section 7.3.1")
-        print(f"  pi* = {[[round(float(z), 4) for z in row] for row in brownian]}, the plan the "
-              f"Gibbs kernel at strength 2 produced")
-
-    potentials = backward_potential(A, g, STEPS)
-    tilted = [tilted_matrix(A, potentials[k], potentials[k + 1]) for k in range(STEPS)]
-    middle = PI_0 @ tilted[0]
-    terminal = middle @ tilted[1]
-    bridge_law = fill_with_reference_bridges(plan, A)
-    bridge_kl = relative_entropy(bridge_law, reference)
-    middle_from_paths = sum(v for k, v in bridge_law.items() if k[1] == 0)
-    if loud:
-        print()
-        for k, h in enumerate(potentials):
-            print(f"h_{k} = ({float(h[0]):.6f}, {float(h[1]):.6f})   "
-                  f"ratio {float(h[1] / h[0]):.2f}")
-        for k, matrix in enumerate(tilted):
-            print(f"A^h_{k} = {[[round(float(z), 6) for z in row] for row in matrix]}")
-        print(f"  (0.5, 0.5) -> {[round(float(z), 6) for z in middle]} -> "
-              f"{[round(float(z), 6) for z in terminal]}")
-        print(f"  1 / h_0 = {[round(float(z), 6) for z in 1.0 / potentials[0]]}, "
-              f"the f the scaling produced")
-        product = torch.diag(PI_0) @ tilted[0] @ tilted[1]  # The endpoint law of the tilted chain.
-        print(f"  diag(pi_0) A^h_0 A^h_1 = "
-              f"{[[round(float(z), 6) for z in row] for row in product]}")
-        print("\npath      R          P*         P*/R")
-        for path in reference:
-            name = ",".join(str(s + 1) for s in path)
-            print(f"  {name}   {reference[path]:.6f}   {bridge_law[path]:.6f}   "
-                  f"{bridge_law[path] / reference[path]:.6f}")
-
-    report = {
-        "R(1,1,2)": reference[(0, 0, 1)],
-        "the eight paths sum to": sum(reference.values()),
-        "chain rule, total": total,
-        "chain rule, marginal term": marginal,
-        "chain rule, conditional term": conditional,
-        "reference terminal error": reference_error,
-        "terminal error, round 1": errors[0],
-        "terminal error, last round": errors[-1],
-        "error ratio per round": reference_error / errors[0],
-        "f": [round(float(z), 6) for z in f],
-        "g": [round(float(z), 6) for z in g],
-        "pi*_11": float(plan[0, 0]),
-        "pi*_12": float(plan[0, 1]),
-        "KL(pi* || R_01)": kl_vectors(plan.flatten(), R01.flatten()),
-        "KL(pi_0 x pi_1 || R_01)": kl_vectors(independent.flatten(), R01.flatten()),
-        "h_0": [round(float(z), 6) for z in potentials[0]],
-        "h_2 = g": [round(float(z), 6) for z in potentials[-1]],
-        "A^h_0(1,1)": float(tilted[0][0, 0]),
-        "A^h_1(1,2)": float(tilted[1][0, 1]),
-        "rows of A^h_0 sum to": float(tilted[0].sum(dim=1).max()),
-        "marginal at the middle step": float(middle[0]),
-        "marginal at the last step": [round(float(z), 6) for z in terminal],
-        "f against 1 / h_0": float((f - 1.0 / potentials[0]).abs().max()),
-        "KL(P* || R) over the paths": bridge_kl,
-        "middle state from the paths": middle_from_paths,
-    }
-    print_report("The static Schrodinger bridge, Chapter 7 Sections 7.12 to 7.15.3", report)
-    save_run(args.out, config=vars(args), report=report)  # Written only with --out.
-
-
-if __name__ == "__main__":
-    main()