Download lecture_7/discrete_bridge.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 13.6 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/discrete_bridge.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_7/discrete_bridge.py
-
curl -L -o discrete_bridge.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/discrete_bridge.py
13.6 kB
| """The Schrodinger bridge over a finite set: rates in place of drifts. | |
| A finite state space has 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, and the two potentials are the Sinkhorn scalings of the reference endpoint law. 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 # Double precision for the matrix exponentials and the time integral. | |
| 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 | |
| # | |
| # We collect the one-step relative entropies and let the step go to zero. What is left is a time | |
| # integral of the Bregman divergence of u log u between the two rates. The divergence is | |
| # nonnegative and vanishes only when the two rates agree. It 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. Sinkhorn | |
| # scales that matrix to the two marginals, and those scalings are the two potentials. | |
| # 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. We | |
| # multiply each off-diagonal rate by the ratio of that potential at the destination to its value at | |
| # the current state, and then set the diagonal 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, a quantity with 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. This 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 | |
| # | |
| # We match the probability flux on an edge in both directions. That fixes the reverse rate at | |
| # 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() | |