Blame for lecture_7/dsbm.py

"""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. The Markov
projection averages that drift over the endpoints of all bridges through a state. That removes
the extra information and leaves a diffusion with the same marginals. The reciprocal projection
replaces the bridges of a Markov law by reference bridges and keeps its endpoint coupling. DSBM
alternates the two, so both marginals stay fixed throughout and only the pairing improves.

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, down to 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.
# If we fit a single number to the two targets by least squares we land 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. To
# project a Markov law we keep its endpoint coupling and redraw 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
#
# We retain the one-step conditionals of a path law and multiply them into a chain. That gives
# the Markov law with the same marginals. 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, fitted by the network 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
#
# A Markov projection leaves every marginal unchanged and a reciprocal projection leaves the
# endpoint coupling unchanged. Both pi_0 and pi_1 therefore 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()