lecture_5/rectify.py

"""Rectification and AReUReDi: changing the coupling so that fewer steps suffice.

A one-step sampler draws every position from its own marginal, so whatever structure the data has
between positions is lost in that draw. The error it makes is the conditional total correlation of
the coupling, a property of how the pairs were drawn and not of the data. We train on the current
pairs, sample from the model, and keep the pairs it produces. That lowers the total correlation. By
the discrete counterpart of the transport-cost bound, it never rises. AReUReDi then composes a
rectified flow with a Metropolis refinement whose proposal starts from the model's own per-position
distribution.

Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).

Chapter 5, Sections 5.36 to 5.41. The second recoupling runs on adjacent base pairs of the synthetic
DNA of ``cis6270.data``. Prints the total correlation of each round, the Hamming distance and the
one-step success rate, and the AReUReDi proposal with its acceptance ratio.

    python lecture_5/rectify.py                     # four rounds, 600 quadrature points each
    python lecture_5/rectify.py --rounds 6          # closer to the limiting coupling
    python lecture_5/rectify.py --steps 20 --quiet  # a coarse integration, visibly less accurate
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.data import DNA, DNA_SIZE, dna_splits
from cis6270.runner import common_parser, print_report, save_run, set_seed


torch.set_default_dtype(torch.double)  # Full precision, because the recoupling runs to within 1e-6 of t = 1.
LENGTH = 2  # Two positions, so the state space has 16 sequences and can be enumerated.
STATES = DNA_SIZE ** LENGTH
LETTER = torch.stack([  # [L, 16] the letter each sequence index holds at each position.
    torch.arange(STATES) // DNA_SIZE, torch.arange(STATES) % DNA_SIZE,
])


# %% 02. The pairing of source and data
#
# Since the first Euler step we have replaced the joint distribution of the next state by the
# product of the per-position marginals. With small steps that costs almost nothing. One large step
# draws every position at once from a product, and the divergence between the joint and that product
# is the conditional total correlation of the coupling.


def position_marginals(row):
    """The per-position marginals of one distribution over sequences. row: [16] -> [L, K]."""
    out = torch.zeros(LENGTH, DNA_SIZE)
    for i in range(LENGTH):
        out[i].index_add_(0, LETTER[i], row)  # Collect the mass of every sequence by its letter.
    return out


def total_correlation(coupling, source):
    """TC_pi(X_1 | X_0), the expected divergence from the joint to the product of its marginals."""
    total = 0.0
    for s in range(STATES):
        marginal = position_marginals(coupling[s])  # [L, K] given this source sequence.
        product = marginal[0][LETTER[0]] * marginal[1][LETTER[1]]  # [16] the factorized draw.
        live = coupling[s] > 0
        total += float(source[s] * (coupling[s][live] * (coupling[s][live] / product[live]).log()).sum())
    return total  # Zero exactly when the positions are independent given the source.


def one_step_success(coupling, source, support):
    """The chance a single factorized draw lands inside the data, averaged over source sequences."""
    total = 0.0
    for s in range(STATES):
        marginal = position_marginals(coupling[s])
        total += float(source[s] * sum(marginal[0][a] * marginal[1][b] for a, b in support))
    return total  # 1/4 under the independent coupling on the complementary pairs.


def hamming(coupling, source):
    """The expected number of positions at which a source and its paired data sequence differ."""
    distance = (LETTER[:, :, None] != LETTER[:, None, :]).sum(dim=0).to(coupling)  # [16, 16].
    return float((source[:, None] * coupling * distance).sum())  # 1.5 under independent pairing.


# %% 03. Total correlation on two bases
#
# Let the data be the four complementary pairs, each with probability one quarter, and the source be
# uniform over all sixteen ordered pairs. Each position of the data alone is uniform over the four
# bases, so the product of the two marginals spreads one sixteenth everywhere while the joint sits on
# four states with one quarter apiece.


def complementary_pairs():
    """The data distribution over AT, TA, CG and GC, and the indices of those four sequences."""
    support = [(DNA.index(a), DNA.index(b)) for a, b in ("AT", "TA", "CG", "GC")]
    q = torch.zeros(STATES)
    for a, b in support:
        q[a * DNA_SIZE + b] = 0.25  # Each complementary pair with probability one quarter.
    return q, support


def adjacent_pairs(tokens):
    """The empirical distribution of adjacent base pairs in the training DNA. Returns [16]."""
    left, right = tokens[:, :-1].reshape(-1), tokens[:, 1:].reshape(-1)  # Every adjacent pair.
    counts = torch.bincount(left * DNA_SIZE + right, minlength=STATES).to(torch.get_default_dtype())
    return counts / counts.sum()  # [16] over the sixteen ordered pairs.


# %% 04. ReDi: one recoupling round
#
# The model here is the entire many-step sampler, and the next coupling pairs a source sample with
# the sequence the sampler produces from it. The velocity is the true marginal velocity of the
# current coupling, and we integrate the master equation without discretization error, so the
# factorization across the two positions is the only approximation left.


def generator(joint, t):
    """The factorized velocity at time t, in the variable u with t = 1 - exp(-u). Returns [16, 16]."""
    eye = torch.eye(DNA_SIZE)
    mix = (1 - t) * eye[:, :, None] + t * eye[:, None, :]  # [K, K, K] the mixture path at one position.
    per = [mix[:, LETTER[i][:, None], LETTER[i][None, :]] for i in range(LENGTH)]  # [K, 16, 16] each.
    weight = joint[None, None, :, :] * per[0][:, None, :, :] * per[1][None, :, :, :]
    weight = weight.reshape(STATES, STATES, STATES)  # [x, x_0, x_1] the unnormalized posterior.
    norm = weight.sum(dim=(1, 2)).clamp(min=1e-300)
    posterior = weight.sum(dim=1) / norm[:, None]  # [x, x_1] the posterior over the data sequence.
    rates, columns = torch.zeros(STATES, STATES), torch.arange(STATES)
    for i in range(LENGTH):
        denoiser = torch.zeros(STATES, DNA_SIZE)  # [x, a] the denoiser at position i.
        denoiser.index_add_(1, LETTER[i], posterior)
        for a in range(DNA_SIZE):
            moved = LETTER.clone()
            moved[i] = a  # [L, 16] every sequence with position i replaced by the letter a.
            destination = moved[0] * DNA_SIZE + moved[1]  # [16] the index of that sequence.
            keep = LETTER[i] != a  # A sequence already holding a at position i sends nothing.
            rates[destination[keep], columns[keep]] += denoiser[keep, a]
    return rates - torch.diag(rates.sum(dim=0))  # In u-time the schedule factor is one.


def recouple(coupling, source, steps, horizon=14.0):
    """Run the exact sampler from every source sequence and return the coupling it produces.

    The substitution t = 1 - exp(-u) holds the schedule factor at one, so the integration stays
    accurate as t approaches the horizon, where 1 - t is about 1e-6. Each cell takes a Heun step,
    second order and accurate to the chapter's figures at a few hundred cells.
    """
    joint = source[:, None] * coupling  # pi(x_0, x_1) over pairs.
    probability = torch.eye(STATES)  # Row s is the state distribution started at the source s.
    du = horizon / steps
    for k in range(steps):
        left = generator(joint, 1.0 - math.exp(-k * du))  # The velocity at the start of the cell.
        right = generator(joint, 1.0 - math.exp(-(k + 1) * du))  # And at its end.
        first = probability @ left.T  # The Euler slope.
        second = (probability + du * first) @ right.T  # The slope the Euler step lands on.
        probability = probability + 0.5 * du * (first + second)  # Their average.
        probability = probability.clamp(min=0.0)
        probability = probability / probability.sum(dim=1, keepdim=True)
    return probability  # [16, 16] the next coupling, one row per source sequence.


# %% 05. AReUReDi: the Tchebycheff score and the tilt
#
# If we score a sequence by its weakest weighted objective, the sampler cannot trade one property
# away. The target is the base model times the exponentiated score, and the refinement anneals the
# strength linearly over its iterations, so early ones explore and late ones commit.


def tchebycheff(scores, weights):
    """S_omega(x) = min_n omega_n s_n(x), the smallest weighted objective. scores: [M, N] -> [M]."""
    return (scores * weights).min(dim=-1).values  # A larger weight puts less pressure on it.


def anneal(iteration, total, low=1.0, high=20.0):
    """eta_t = eta_min + (eta_max - eta_min) t/(T - 1), or 1 + t at the chapter's settings."""
    return low + (high - low) * iteration / (total - 1)  # 12 at the twelfth iteration of twenty.


# %% 06. The proposal and its acceptance step
#
# The proposal starts from the model's own per-position distribution and reweights it by a balancing
# function of how much the tilt would change. The identity g(u) = u g(1/u) collapses the acceptance
# ratio to a ratio of two one-position normalizers, at the cost of one sum over the alphabet.


def barker(u):
    """g(u) = u/(1 + u), and g(u) = u g(1/u)."""
    return u / (1.0 + u)  # The square root is the other common choice.


def proposal(model_distribution, tilts, current):
    """q_i(. | x) and its normalizer Z(x). tilts: [K] the value of W at each letter."""
    ratios = tilts / tilts[current]  # r_i(y; x), the tilt at the proposed letter over the current one.
    weighted = model_distribution * barker(ratios)  # The model's prior times the balancing function.
    return weighted / weighted.sum(), float(weighted.sum()), ratios  # q, Z(x) and the ratios.


def acceptance(z_current, z_proposed):
    """alpha_i(x, x') = min(1, Z(x)/Z(x')), what the balancing identity leaves of the ratio."""
    return min(1.0, z_current / z_proposed)  # One in both directions when the normalizers agree.


# %% 07. Concentration on the maximizers
#
# If we keep only the maximizers in the denominator and bound the numerator by the score gap, we get
# an exponential bound on the probability of anything outside the maximizing set, so the samples
# concentrate there as the strength grows.


def tilted_distribution(base, scores, eta):
    """pi_{eta, omega}(x) proportional to p_1(x) exp(eta S_omega(x))."""
    weighted = base * (eta * scores).exp()
    return weighted / weighted.sum()  # [M] over the sequences under consideration.


def concentration_bound(base, scores, eta):
    """exp(-eta Delta) p_1(x) over the base mass of the maximizers, for a sequence outside the set."""
    best = scores.max()
    gap = float(best - scores[scores < best].max())  # Delta_omega, the gap to the next best score.
    return math.exp(-eta * gap) / float(base[scores == best].sum()), gap


# %% 08. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--rounds", type=int, default=4, help="rectification rounds")
    args = parser.parse_args()
    set_seed(args.seed)

    source = torch.full((STATES,), 1.0 / STATES)  # Uniform over all sixteen ordered pairs.
    q, support = complementary_pairs()
    coupling = q[None, :].repeat(STATES, 1)  # The independent coupling pairs every source with q.
    report = {
        "tc_round_0": round(total_correlation(coupling, source), 3),  # log 4 = 1.386.
        "log_4": round(math.log(4), 3),  # The enumeration of the four pairs gives this.
        "hamming_round_0": round(hamming(coupling, source), 3),  # 1.5.
        "success_round_0": round(one_step_success(coupling, source, support), 4),  # 0.25.
    }
    for round_index in range(1, args.rounds + 1):
        coupling = recouple(coupling, source, args.steps)  # One application of the operator.
        report[f"tc_round_{round_index}"] = round(total_correlation(coupling, source), 3)
        report[f"hamming_round_{round_index}"] = round(hamming(coupling, source), 3)
        report[f"success_round_{round_index}"] = round(one_step_success(coupling, source, support), 4)
        if round_index == 1:  # Where the first round's gain comes from, read at two sources.
            columns = [a * DNA_SIZE + b for a, b in support]  # The four complementary pairs.
            for name in ("AA", "AT"):  # One source outside the data and one inside it.
                row = DNA.index(name[0]) * DNA_SIZE + DNA.index(name[1])
                report[f"pairing_from_{name}"] = [round(float(coupling[row, c]), 3) for c in columns]
    report["tc_limit"] = round(0.75 * math.log(2), 3)  # Twelve sources split in half, 0.520 nats.
    report["success_limit"] = round(10.0 / 16.0, 4)  # 0.625, four exact sources and twelve halves.
    report["hamming_limit"] = round(12.0 / 16.0, 4)  # 0.75, the least any coupling can reach.

    # The same measurement on adjacent base pairs of the synthetic DNA, under independent pairing.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    dna = adjacent_pairs(splits["train_tokens"])  # [16] the empirical distribution of a base pair.
    dna_coupling = dna[None, :].repeat(STATES, 1)  # Every source paired with the same distribution.
    live = dna[dna > 0]
    report["dna_pair_entropy"] = round(float(-(live * live.log()).sum()), 4)  # At most log 16 = 2.77.
    report["dna_total_correlation"] = round(total_correlation(dna_coupling, source), 4)
    report["dna_fraction_of_entropy"] = round(
        total_correlation(dna_coupling, source) / float(-(live * live.log()).sum()), 4)

    # Section 5.41.2: one refinement iteration at the twelfth of twenty, with two objectives.
    scores = torch.tensor([[0.60, 0.80], [0.70, 0.76], [0.50, 0.90], [0.64, 0.62]])
    weights = torch.tensor([0.5, 0.5])
    eta = anneal(11, 20)  # The twelfth iteration, counting from zero, where 1 + t equals 12.
    tchebycheff_scores = tchebycheff(scores, weights)
    tilts = (eta * tchebycheff_scores).exp()  # W = exp(eta S) at that strength.
    report["annealed_eta"] = round(eta, 4)  # 1 + t over twenty iterations.
    report["tchebycheff"] = [round(float(s), 3) for s in tchebycheff_scores]  # (0.300, ...).
    report["tilts"] = [round(float(w), 2) for w in tilts]  # (36.60, 66.69, 20.09, 41.26).
    model_distribution = torch.tensor([0.40, 0.25, 0.20, 0.15])  # The rectified model at this position.
    q_current, z_current, ratios = proposal(model_distribution, tilts, 0)  # The position holds A.
    report["ratios_from_A"] = [round(float(r), 3) for r in ratios]  # (1.000, 1.822, 0.549, 1.127).
    report["barker_from_A"] = [round(float(barker(r)), 3) for r in ratios]  # (0.500, 0.646, ...).
    report["normalizer_at_A"] = round(z_current, 4)  # 0.5118.
    report["proposal_at_A"] = [round(float(p), 3) for p in q_current]  # (0.391, 0.315, 0.138, 0.155).
    _, z_proposed, _ = proposal(model_distribution, tilts, 1)  # The proposed sequence holds C.
    report["normalizer_at_C"] = round(z_proposed, 4)  # 0.3704.
    report["accept_A_to_C"] = round(acceptance(z_current, z_proposed), 3)  # One.
    report["accept_C_to_A"] = round(acceptance(z_proposed, z_current), 3)  # 0.724.

    # Section 5.41.3: the tilted distribution over four sequences, and the bound on a non-maximizer.
    base = torch.tensor([0.50, 0.25, 0.15, 0.10])
    objective = torch.tensor([0.20, 0.30, 0.45, 0.45])
    for strength in (20.0, 60.0):
        tilted = tilted_distribution(base, objective, strength)
        bound, gap = concentration_bound(base, objective, strength)
        report[f"mass_on_maximizers_eta_{strength:g}"] = f"{float(tilted[2:].sum()):.6f}"
        report[f"second_sequence_eta_{strength:g}"] = f"{float(tilted[1]):.6f}"
        report[f"bound_eta_{strength:g}"] = f"{bound * float(base[1]):.6f}"
    report["score_gap"] = round(gap, 4)  # 0.15.
    print_report("rectification", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":
    main()