"""Blind recovery on synthetic set-kernel fields: the end-to-end test. Builds the connected-component descriptor field and the phrase bag-of-words field, hides the text order behind a shuffle, and runs parallel tempering on the relational energy alone. Recovery accuracy against the hidden truth is the first end-to-end measurement of world matching in the closed world. """ from __future__ import annotations import argparse import json from pathlib import Path import torch from tqdm import tqdm from .blind_recovery import arm_tempering, permutation_energy_batch from .common import read_json, seed_everything, write_json from .manifold_gate import standardize_relation from .synth_cc_battery import ( component_descriptors, moment_field, onehot_descriptors, phrase_bow_sets, ) from .synth_set_battery import set_similarity_field 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("--replicas", type=int, default=8) parser.add_argument("--tempering-rounds", type=int, default=60000) 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("--unary-weight", type=float, default=0.0) parser.add_argument("--init", default="random") parser.add_argument("--residualize-size", action="store_true", default=False) parser.add_argument( "--features", choices=["descriptors", "onehot"], default="onehot" ) parser.add_argument("--kernel", choices=["matching", "moment"], default="moment") parser.add_argument("--device", default="cuda:3") parser.add_argument("--seed", type=int, default=20260731) parser.add_argument( "--output", default="artifacts/synth_v0/recovery_end_to_end.json" ) return parser.parse_args() def residualize(field: torch.Tensor, sizes: torch.Tensor) -> torch.Tensor: """Regress the set-size nuisance out of a matching-value field. Set sizes are unimodal observables; their sum, difference, and product explain a size-driven component that differs between modalities and is exploitable by counterfeit assignments. """ n = len(field) mask = ~torch.eye(n, dtype=torch.bool) features = torch.stack( [ (sizes[:, None] + sizes[None, :])[mask], (sizes[:, None] - sizes[None, :]).abs()[mask], (sizes[:, None] * sizes[None, :])[mask], torch.ones(int(mask.sum()), dtype=torch.float64), ], dim=1, ) values = field.double()[mask] solution = torch.linalg.lstsq(features, values[:, None]).solution residual = values - (features @ solution).squeeze(1) output = field.double().clone() output[mask] = residual output.fill_diagonal_(0.0) return output.float() def main() -> None: args = parse_args() seed_everything(args.seed) manifest = read_json(Path(args.data_dir, "manifest.json")) captions = read_json(Path(args.data_dir, "captions.json"))["captions"] rows = manifest[args.split][: args.samples] image_dir = Path(manifest["image_dir"]) import torch.nn.functional as F per_view_fields = [] for view in range(args.vision_views): raw_sets = [] for row in tqdm(rows, desc=f"cc v{view}"): image = load_image(image_dir / f"scene{row:06d}_v{view}.png") sprites, _ = component_descriptors(image, args.merge_distance) raw_sets.append(sprites) if args.features == "onehot": vision_sets = [ F.normalize(v, dim=-1) for v in onehot_descriptors(raw_sets) ] else: vision_sets = [F.normalize(v, dim=-1) for v in raw_sets] if args.kernel == "moment": per_view_fields.append(moment_field(vision_sets)) else: per_view_fields.append(set_similarity_field(vision_sets)) visual_field = torch.stack(per_view_fields).mean(0) text_sets = phrase_bow_sets(rows, captions, manifest["vocabulary"]) if args.kernel == "moment": text_field = moment_field(text_sets) else: text_field = set_similarity_field(text_sets) if args.residualize_size: vision_sizes = torch.tensor( [len(s) for s in vision_sets], dtype=torch.float64 ) text_sizes = torch.tensor([len(s) for s in text_sets], dtype=torch.float64) visual_field = residualize(visual_field, vision_sizes) text_field = residualize(text_field, text_sizes) 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) text_input = text_field[hidden][:, hidden] text_standardized = standardize_relation(text_input.double())[0].float().to(device) visual_standardized = ( standardize_relation(visual_field.double())[0].float().to(device) ) true_energy = float( permutation_energy_batch( text_standardized, visual_standardized, truth[None].to(device) )[0] ) report = { "protocol": ( "Set-kernel fields from released renders and captions; text " "order hidden behind a shuffle; tempering sees no truth. " "Hidden truth scores the outcome only." ), "split": args.split, "samples": size, "true_energy": true_energy, "chance_accuracy": 1.0 / size, "tempering": arm_tempering( text_standardized, visual_standardized, truth, true_energy, args, generator, unary=None, ), } best = max( report["tempering"]["replicas"] + [report["tempering"]["best"]], key=lambda c: c["accuracy"], ) by_energy = min( report["tempering"]["replicas"] + [report["tempering"]["best"]], key=lambda c: c["energy"], ) report["summary"] = { "best_accuracy": best["accuracy"], "best_accuracy_energy_over_true": best["energy_over_true"], "lowest_energy_accuracy": by_energy["accuracy"], "lowest_energy_over_true": by_energy["energy_over_true"], } print(json.dumps({"summary": report["summary"]})) write_json(args.output, report) print(f"Wrote {args.output}") if __name__ == "__main__": main()