path: root/worldalign/synth_fast_gate.py

"""Closed-form batched gate: exact all-triple triangle energy.

The sampled triangle energy cost 5.7 ms per proposal because each
evaluation gathered a variable-length triple set in Python. It has a
closed form. Writing M(sigma) for the elementwise product of the
sigma-permuted text field with the visual field,

    pairwise    = const - 2 * sum(M)
    all-triple  = const - 2 * trace(M^3) / 6

because trace of a power is invariant under simultaneous row-column
permutation, so the text-only and vision-only terms do not move. One
batched matrix product evaluates trace(M^3) = sum(M * (M @ M)) for
hundreds of candidate permutations at once, over every C(N,3) triple
rather than a sample.

This buys a genuinely stronger searcher: exact steepest descent over all
N(N-1)/2 transpositions per step, plus batched-proposal tempering.
Hidden pairs score orderings only.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch

from .common import read_json, seed_everything, write_json
from .synth_triangle_gate import build_fields, standardized


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", default="", help="Saved .pt with visual_field/text_field.")
    parser.add_argument("--data-dir", default="artifacts/synth_v0")
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=512)
    parser.add_argument("--merge-distance", type=float, default=30.0)
    parser.add_argument("--vision-views", type=int, default=4)
    parser.add_argument("--energies", default="pair,triangle,both")
    parser.add_argument("--chunk", type=int, default=512)
    parser.add_argument("--descent-restarts", type=int, default=5)
    parser.add_argument("--descent-max-steps", type=int, default=4000)
    parser.add_argument("--replicas", type=int, default=8)
    parser.add_argument("--rounds", type=int, default=3000)
    parser.add_argument("--proposals", type=int, default=64)
    parser.add_argument("--temp-high", type=float, default=3e-2)
    parser.add_argument("--temp-low", type=float, default=1e-4)
    parser.add_argument("--exchange-every", type=int, default=20)
    parser.add_argument("--polish-steps", type=int, default=2000)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--seed", type=int, default=20260731)
    parser.add_argument("--output", default="artifacts/synth_v0/fast_gate.json")
    return parser.parse_args()


class ClosedFormEnergy:
    """Energy as a function of M = permuted-text * visual, batched."""

    def __init__(
        self,
        text: torch.Tensor,
        visual: torch.Tensor,
        pair_weight: float,
        triangle_weight: float,
        chunk: int,
    ) -> None:
        self.text = text
        self.visual = visual
        self.pair_weight = pair_weight
        self.triangle_weight = triangle_weight
        self.chunk = chunk
        self.size = len(visual)
        self.pair_count = self.size * (self.size - 1)
        self.triple_count = self.size * (self.size - 1) * (self.size - 2)
        # Permutation-invariant constants, kept so reported values match
        # the direct definitions of the two energies.
        square_text = text * text
        square_visual = visual * visual
        self.pair_constant = float(square_text.sum() + square_visual.sum())
        self.triangle_constant = float(
            self._trace_cube(square_text[None])[0]
            + self._trace_cube(square_visual[None])[0]
        )

    @staticmethod
    def _trace_cube(matrices: torch.Tensor) -> torch.Tensor:
        return (matrices * torch.bmm(matrices, matrices)).sum((-2, -1))

    def energy(self, permutations: torch.Tensor) -> torch.Tensor:
        """Exact energy for a batch of permutations [B, N]."""
        values = []
        for start in range(0, len(permutations), self.chunk):
            block = permutations[start : start + self.chunk]
            permuted = self.text[block[:, :, None], block[:, None, :]]
            product = permuted * self.visual
            total = torch.zeros(len(block), device=permuted.device)
            if self.pair_weight:
                pair = (self.pair_constant - 2.0 * product.sum((-2, -1))) / (
                    self.pair_count
                )
                total = total + self.pair_weight * pair
            if self.triangle_weight:
                triangle = (
                    self.triangle_constant - 2.0 * self._trace_cube(product)
                ) / self.triple_count
                total = total + self.triangle_weight * triangle
            values.append(total)
        return torch.cat(values)


def all_swaps(size: int, device: torch.device) -> torch.Tensor:
    rows, cols = torch.triu_indices(size, size, offset=1, device=device)
    return torch.stack([rows, cols], dim=1)


def apply_swaps(permutation: torch.Tensor, swaps: torch.Tensor) -> torch.Tensor:
    batch = permutation[None].repeat(len(swaps), 1)
    index = torch.arange(len(swaps), device=permutation.device)
    p, q = swaps[:, 0], swaps[:, 1]
    values_p = batch[index, p].clone()
    batch[index, p] = batch[index, q]
    batch[index, q] = values_p
    return batch


def steepest_descent(
    energy: ClosedFormEnergy,
    start: torch.Tensor,
    swaps: torch.Tensor,
    max_steps: int,
) -> tuple[torch.Tensor, float]:
    """Exact steepest descent over every transposition each step."""
    current = start.clone()
    value = float(energy.energy(current[None])[0])
    for _ in range(max_steps):
        candidates = apply_swaps(current, swaps)
        values = energy.energy(candidates)
        best = int(values.argmin())
        if float(values[best]) >= value - 1e-12:
            break
        current = candidates[best]
        value = float(values[best])
    return current, value


def temper(
    energy: ClosedFormEnergy,
    starts: torch.Tensor,
    args: argparse.Namespace,
    generator: torch.Generator,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Batched-proposal parallel tempering; returns states and energies."""
    size = energy.size
    replicas = len(starts)
    temperatures = torch.logspace(
        torch.log10(torch.tensor(args.temp_low)),
        torch.log10(torch.tensor(args.temp_high)),
        replicas,
    ).to(device)
    states = starts.clone()
    values = energy.energy(states)
    for round_index in range(args.rounds):
        p = torch.randint(
            0, size, (replicas, args.proposals), generator=generator
        ).to(device)
        q = torch.randint(
            0, size, (replicas, args.proposals), generator=generator
        ).to(device)
        valid = p != q
        batch = states[:, None, :].repeat(1, args.proposals, 1)
        index_r = torch.arange(replicas, device=device)[:, None]
        original_p = batch.gather(2, p[..., None]).squeeze(-1)
        original_q = batch.gather(2, q[..., None]).squeeze(-1)
        batch.scatter_(2, p[..., None], original_q[..., None])
        batch.scatter_(2, q[..., None], original_p[..., None])
        flat = batch.reshape(replicas * args.proposals, size)
        proposal_values = energy.energy(flat).reshape(replicas, args.proposals)
        deltas = proposal_values - values[:, None]
        noise = torch.rand(
            replicas, args.proposals, generator=generator
        ).to(device)
        threshold = -temperatures[:, None] * noise.clamp_min(1e-12).log()
        accept = (deltas < threshold) & valid
        first = torch.where(
            accept.any(-1),
            accept.float().argmax(-1),
            torch.zeros(replicas, dtype=torch.long, device=device),
        )
        taken = accept.any(-1)
        chosen = batch[index_r.squeeze(-1), first]
        states = torch.where(taken[:, None], chosen, states)
        values = torch.where(
            taken, proposal_values[index_r.squeeze(-1), first], values
        )
        if round_index % args.exchange_every == 0:
            for replica in range(replicas - 1):
                gap = (values[replica] - values[replica + 1]) * (
                    1.0 / temperatures[replica] - 1.0 / temperatures[replica + 1]
                )
                accept_swap = gap > 0 or float(
                    torch.rand(1, generator=generator)
                ) < float(gap.exp().clamp(max=1.0))
                if accept_swap:
                    states[[replica, replica + 1]] = states[
                        [replica + 1, replica]
                    ]
                    values[[replica, replica + 1]] = values[
                        [replica + 1, replica]
                    ]
    return states, values


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    if args.fields:
        state = torch.load(args.fields, map_location="cpu", weights_only=False)
        visual_field, text_field = state["visual_field"], state["text_field"]
    else:
        manifest = read_json(Path(args.data_dir, "manifest.json"))
        rows = manifest[args.split][: args.samples]
        visual_field, text_field = build_fields(args, rows, manifest)

    device = torch.device(args.device)
    size = len(visual_field)
    generator = torch.Generator().manual_seed(args.seed)
    hidden = torch.randperm(size, generator=generator)
    truth = torch.argsort(hidden).to(device)  # scored, never optimized against
    text = standardized(text_field[hidden][:, hidden].double().to(device)).float()
    visual = standardized(visual_field.double().to(device)).float()
    swaps = all_swaps(size, device)

    weights = {"pair": (1.0, 0.0), "triangle": (0.0, 1.0), "both": (1.0, 1.0)}
    report = {
        "protocol": (
            "Closed-form energies over all triples; the decision statistic "
            "is the energy of the truth against the deepest state reached "
            "by exact steepest descent and batched tempering. Hidden pairs "
            "score only."
        ),
        "samples": size,
        "energies": {},
    }
    for name in (item.strip() for item in args.energies.split(",")):
        pair_weight, triangle_weight = weights[name]
        energy = ClosedFormEnergy(
            text, visual, pair_weight, triangle_weight, args.chunk
        )
        true_energy = float(energy.energy(truth[None])[0])
        random_batch = torch.stack(
            [
                torch.argsort(torch.rand(size, generator=generator))
                for _ in range(200)
            ]
        ).to(device)
        random_values = energy.energy(random_batch)

        kept, kept_value = steepest_descent(
            energy, truth, swaps, args.descent_max_steps
        )
        retention = float((kept.cpu() == truth.cpu()).float().mean())

        quenches = []
        for restart in range(args.descent_restarts):
            start = torch.argsort(torch.rand(size, generator=generator)).to(device)
            final, value = steepest_descent(
                energy, start, swaps, args.descent_max_steps
            )
            quenches.append(
                {
                    "energy": value,
                    "accuracy": float((final.cpu() == truth.cpu()).float().mean()),
                }
            )

        starts = torch.stack(
            [
                torch.argsort(torch.rand(size, generator=generator))
                for _ in range(args.replicas)
            ]
        ).to(device)
        states, values = temper(energy, starts, args, generator, device)
        cold = int(values.argmin())
        polished, polished_value = steepest_descent(
            energy, states[cold], swaps, args.polish_steps
        )
        tempering = {
            "cold_energy": float(values[cold]),
            "polished_energy": polished_value,
            "polished_accuracy": float(
                (polished.cpu() == truth.cpu()).float().mean()
            ),
            "best_replica_accuracy": max(
                float((state.cpu() == truth.cpu()).float().mean())
                for state in states
            ),
        }
        deepest = min(
            [polished_value, float(values.min())]
            + [item["energy"] for item in quenches]
        )
        entry = {
            "true_energy": true_energy,
            "random_mean": float(random_values.mean()),
            "true_z": float(
                (random_values.mean() - true_energy)
                / random_values.std().clamp_min(1e-12)
            ),
            "descent_retention_diagnostic": retention,
            "descent_from_truth_energy": kept_value,
            "quenches": quenches,
            "tempering": tempering,
            "deepest_seen": deepest,
            "margin_over_true": deepest / abs(true_energy) - true_energy / abs(true_energy),
            "passes": bool(deepest >= true_energy - 1e-9),
            "recovery_accuracy": tempering["polished_accuracy"],
        }
        report["energies"][name] = entry
        print(
            json.dumps(
                {
                    name: {
                        key: entry[key]
                        for key in (
                            "true_energy",
                            "deepest_seen",
                            "passes",
                            "recovery_accuracy",
                            "descent_retention_diagnostic",
                            "true_z",
                        )
                    }
                }
            )
        )
    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()


def all_pair_swap_deltas(
    permuted_text: torch.Tensor, visual: torch.Tensor
) -> torch.Tensor:
    """Energy change of every transposition at once, for the pairwise term.

    The brute-force descent forms all N(N-1)/2 candidate permutations and
    scores each through the full energy, which at 256 scenes means 32,640
    matrices per step and, with the triangle term active, a batched cube
    trace over every one of them -- 203 seconds for a single descent. The
    pairwise term does not need any of that. Writing S for the alignment sum
    and C for `permuted_text @ visual`, swapping positions p and q changes S
    by 2(C_pq + C_qp - C_pp - C_qq + 2 A_pq B_pq), so one matrix product
    yields the entire table.

    Both fields must be symmetric with zero diagonal, which standardisation
    already guarantees. Returns a dense [N, N] table of energy deltas; the
    sign convention matches `ClosedFormEnergy.energy`, so negative is an
    improvement.
    """
    size = permuted_text.shape[-1]
    product = permuted_text @ visual
    diagonal = torch.diagonal(product)
    gain = 2.0 * (
        product + product.T
        - diagonal[:, None] - diagonal[None, :]
        + 2.0 * permuted_text * visual
    )
    return -2.0 * gain / (size * (size - 1))


def fast_pair_descent(
    text: torch.Tensor, visual: torch.Tensor, start: torch.Tensor, max_steps: int
) -> torch.Tensor:
    """Steepest descent on the pairwise energy using the closed-form table."""
    size = len(visual)
    current = start.clone()
    triangle = torch.triu(
        torch.ones(size, size, dtype=torch.bool, device=visual.device), diagonal=1
    )
    for _ in range(max_steps):
        permuted = text[current[:, None], current[None, :]]
        deltas = all_pair_swap_deltas(permuted, visual)
        deltas = torch.where(triangle, deltas, torch.inf)
        best = int(deltas.argmin())
        if float(deltas.view(-1)[best]) >= -1e-12:
            break
        p, q = best // size, best % size
        current[[p, q]] = current[[q, p]]
    return current