path: root/worldalign/synth_triangle_gate.py

"""Third-order (triangle) relational invariants for the synthetic world.

Counterfeits so far satisfy pairwise statistics: they are wrong sections
that look flat when probed along edges. The holonomy analogue on a
relation field is a triangle: for a node triple the product-like
statistic T[a,b,c] combines three edges at once, so preserving pairwise
marginals no longer suffices. Each node participates in O(N^2) triangles
rather than O(N) edges, which multiplies the constraint count without
touching the states.

The triangle energy sums squared cross-modal differences of standardized
triangle tensors over a fixed random triple sample; the sample is drawn
once from node indices only. Swap deltas are evaluated exactly on the
triples touching the swapped nodes. Hidden pairs score orderings only.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
import torch.nn.functional as F
from tqdm import tqdm

from .common import read_json, seed_everything, write_json
from .synth_cc_battery import (
    component_descriptors,
    moment_field,
    onehot_descriptors,
    phrase_bow_sets,
)
from .synth_towers import load_image


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    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("--triples", type=int, default=200000)
    parser.add_argument("--pair-weight", type=float, default=1.0)
    parser.add_argument("--triangle-weight", type=float, default=1.0)
    parser.add_argument("--random-perms", type=int, default=200)
    parser.add_argument("--transposition-samples", type=int, default=20000)
    parser.add_argument("--descent-restarts", type=int, default=3)
    parser.add_argument("--descent-max-steps", type=int, default=3000)
    parser.add_argument("--descent-proposals", type=int, default=2048)
    parser.add_argument("--replicas", type=int, default=8)
    parser.add_argument("--tempering-rounds", type=int, default=20000)
    parser.add_argument("--temp-high", type=float, default=3e-3)
    parser.add_argument("--temp-low", type=float, default=1e-5)
    parser.add_argument("--exchange-every", type=int, default=20)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--seed", type=int, default=20260731)
    parser.add_argument(
        "--output", default="artifacts/synth_v0/triangle_gate.json"
    )
    return parser.parse_args()


def standardized(field: torch.Tensor) -> torch.Tensor:
    mask = ~torch.eye(len(field), dtype=torch.bool, device=field.device)
    values = field[mask]
    out = (field - values.mean()) / values.std().clamp_min(1e-9)
    return out.masked_fill(~mask, 0.0)


class TriangleEnergy:
    """Pairwise-plus-triangle energy over a fixed triple sample."""

    def __init__(
        self,
        text: torch.Tensor,
        visual: torch.Tensor,
        triples: torch.Tensor,
        pair_weight: float,
        triangle_weight: float,
    ) -> None:
        self.text = text
        self.visual = visual
        self.triples = triples
        self.pair_weight = pair_weight
        self.triangle_weight = triangle_weight
        self.size = len(visual)
        self.mask = ~torch.eye(self.size, dtype=torch.bool, device=visual.device)
        self.pair_count = int(self.mask.sum())
        a, b, c = triples[:, 0], triples[:, 1], triples[:, 2]
        self.visual_triangle = (
            visual[a, b] * visual[b, c] * visual[a, c]
        )
        # Triples touching each node, for local delta evaluation.
        self.touch = [[] for _ in range(self.size)]
        for index, (x, y, z) in enumerate(triples.tolist()):
            self.touch[x].append(index)
            self.touch[y].append(index)
            self.touch[z].append(index)
        self.touch = [
            torch.tensor(items, device=visual.device, dtype=torch.long)
            for items in self.touch
        ]

    def pairwise(self, permutation: torch.Tensor) -> torch.Tensor:
        fields = self.text[permutation][:, permutation]
        return ((fields - self.visual) ** 2)[self.mask].sum() / self.pair_count

    def triangle(self, permutation: torch.Tensor) -> torch.Tensor:
        fields = self.text[permutation][:, permutation]
        a, b, c = self.triples[:, 0], self.triples[:, 1], self.triples[:, 2]
        text_triangle = fields[a, b] * fields[b, c] * fields[a, c]
        return ((text_triangle - self.visual_triangle) ** 2).mean()

    def total(self, permutation: torch.Tensor) -> float:
        return float(
            self.pair_weight * self.pairwise(permutation)
            + self.triangle_weight * self.triangle(permutation)
        )

    def swap_delta(self, permutation: torch.Tensor, p: int, q: int) -> float:
        """Exact delta for swapping positions p and q."""
        before_pair = self._local_pair(permutation, p, q)
        indices = torch.cat([self.touch[p], self.touch[q]]).unique()
        before_triangle = self._local_triangle(permutation, indices)
        trial = permutation.clone()
        trial[[p, q]] = trial[[q, p]]
        after_pair = self._local_pair(trial, p, q)
        after_triangle = self._local_triangle(trial, indices)
        return float(
            self.pair_weight * (after_pair - before_pair) / self.pair_count
            + self.triangle_weight
            * (after_triangle - before_triangle)
            / len(self.triples)
        )

    def _local_pair(
        self, permutation: torch.Tensor, p: int, q: int
    ) -> torch.Tensor:
        rows = permutation[[p, q]]
        fields = self.text[rows][:, permutation]
        visual_rows = self.visual[[p, q]]
        difference = (fields - visual_rows) ** 2
        difference[0, p] = 0.0
        difference[1, q] = 0.0
        # Rows and columns are symmetric; count row contributions twice and
        # subtract the doubly counted (p, q) pair once.
        total = 2.0 * difference.sum()
        return total - 2.0 * difference[0, q]

    def _local_triangle(
        self, permutation: torch.Tensor, indices: torch.Tensor
    ) -> torch.Tensor:
        triples = self.triples[indices]
        a, b, c = triples[:, 0], triples[:, 1], triples[:, 2]
        pa, pb, pc = permutation[a], permutation[b], permutation[c]
        text_triangle = (
            self.text[pa, pb] * self.text[pb, pc] * self.text[pa, pc]
        )
        return ((text_triangle - self.visual_triangle[indices]) ** 2).sum()


def build_fields(args: argparse.Namespace, rows: list[int], manifest: dict) -> tuple:
    image_dir = Path(manifest["image_dir"])
    per_view = []
    for view in range(args.vision_views):
        raw = []
        for row in tqdm(rows, desc=f"cc v{view}"):
            sprites, _ = component_descriptors(
                load_image(image_dir / f"scene{row:06d}_v{view}.png"),
                args.merge_distance,
            )
            raw.append(sprites)
        sets = [F.normalize(v, dim=-1) for v in onehot_descriptors(raw)]
        per_view.append(moment_field(sets))
    visual_field = torch.stack(per_view).mean(0)
    captions = read_json(Path(args.data_dir, "captions.json"))["captions"]
    text_sets = phrase_bow_sets(rows, captions, manifest["vocabulary"])
    return visual_field, moment_field(text_sets)


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    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(rows)
    generator = torch.Generator().manual_seed(args.seed)
    hidden = torch.randperm(size, generator=generator)
    truth = torch.argsort(hidden).to(device)
    text = standardized(text_field[hidden][:, hidden].double().to(device)).float()
    visual = standardized(visual_field.double().to(device)).float()

    triples = torch.randint(0, size, (args.triples, 3), generator=generator)
    triples = triples[
        (triples[:, 0] != triples[:, 1])
        & (triples[:, 1] != triples[:, 2])
        & (triples[:, 0] != triples[:, 2])
    ].to(device)
    energy = TriangleEnergy(
        text, visual, triples, args.pair_weight, args.triangle_weight
    )

    true_total = energy.total(truth)
    randoms = []
    for _ in range(args.random_perms):
        permutation = torch.argsort(torch.rand(size, generator=generator)).to(device)
        randoms.append(energy.total(permutation))
    randoms = torch.tensor(randoms)
    gate_a = {
        "true": true_total,
        "random_mean": float(randoms.mean()),
        "true_z": float((randoms.mean() - true_total) / randoms.std().clamp_min(1e-12)),
    }

    improving = 0
    checked = 0
    for _ in range(min(args.transposition_samples, 4000)):
        p = int(torch.randint(0, size, (1,), generator=generator))
        q = int(torch.randint(0, size, (1,), generator=generator))
        if p == q:
            continue
        checked += 1
        improving += energy.swap_delta(truth, p, q) < 0
    gate_b = {
        "sampled": checked,
        "improving_fraction": improving / max(checked, 1),
    }

    def descent(start: torch.Tensor) -> dict:
        current = start.clone()
        for _ in range(args.descent_max_steps):
            best_delta, best_pair = 0.0, None
            for _ in range(64):
                p = int(torch.randint(0, size, (1,), generator=generator))
                q = int(torch.randint(0, size, (1,), generator=generator))
                if p == q:
                    continue
                delta = energy.swap_delta(current, p, q)
                if delta < best_delta:
                    best_delta, best_pair = delta, (p, q)
            if best_pair is None:
                break
            p, q = best_pair
            current[[p, q]] = current[[q, p]]
        return {
            "accuracy": float((current.cpu() == truth.cpu()).float().mean()),
            "energy": energy.total(current),
        }

    from_true = descent(truth)
    restarts = [
        descent(torch.argsort(torch.rand(size, generator=generator)).to(device))
        for _ in range(args.descent_restarts)
    ]
    best_random = min(r["energy"] for r in restarts)

    report = {
        "protocol": (
            "Pairwise-plus-triangle energy on moment set-kernel fields; "
            "triples sampled from node indices only. Hidden pairs score "
            "orderings."
        ),
        "samples": size,
        "triples": len(triples),
        "weights": {
            "pair": args.pair_weight,
            "triangle": args.triangle_weight,
        },
        "gate_a": gate_a,
        "gate_b": gate_b,
        "descent_from_true": from_true,
        "descent_from_random": restarts,
        "verdict": {
            "true_energy": true_total,
            "best_random_descent": best_random,
            "counterfeit_found": bool(
                best_random < true_total
                and min(r["accuracy"] for r in restarts) < 0.5
            ),
            "descent_keeps": from_true["accuracy"],
            "true_z": gate_a["true_z"],
            "improving_fraction": gate_b["improving_fraction"],
        },
    }
    print(json.dumps({"verdict": report["verdict"]}))
    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()