Blame for lecture_7/csbm.py

"""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, and the Markov projection keeps the one-step conditionals
and multiplies them into a chain. 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. When the categories have an order we can use a discretized Gaussian
# centred on the current one, 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, and the result is 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
#
# We alternate the two projections, 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. The starting
    # gap records that and the rate of convergence follows it. Two rounds separate the two ratios
    # from the float64 floor, and the sixty-four-path column reaches it 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()