summaryrefslogtreecommitdiff
path: root/worldalign/anchor_bound.py
blob: 49f4ed35ce8afafb64027bbee083a1d2e8038fa1 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
"""A cheap gate that survives its controls, built as an upper bound.

Three statistics have now been proposed as go/no-go predictors and three have
failed: the field correlation, the width of the shared spectrum, and the
held-out shared dimension. Each failed the same way -- fields agreeing on the
statistic and disagreeing on recovery. The lesson is that a predictor invented
by staring at the fields keeps finding quantities that are not about matching.

So this asks the matching question directly, with one thing given away. Half
the scenes are declared anchors and their correspondence is handed over. Every
remaining scene is described, in each modality, by its field row against the
anchors, and the two descriptions are matched one-to-one by Hungarian
assignment. Nothing is searched: the answer is read off a linear program.

The result upper-bounds what blind recovery can achieve, because blind recovery
must discover the anchor correspondence as well. It is therefore a **necessary
condition**: a representation that cannot be matched with the anchors given
will not be matched without them, and it can be rejected in seconds rather than
after a search. It is an empirical bound rather than a proven one -- the given
anchors also change the estimation problem, not only remove search -- but it
held in every field measured, and unlike the three refuted statistics it
separates the cases they collapse.

It also splits the deficit in two, which is what makes it diagnostic rather
than merely predictive. The bound itself measures whether the representation
identifies scenes at all. The gap between the bound and blind recovery measures
whether the search can find what the representation contains.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch
from scipy.optimize import linear_sum_assignment

from .common import write_json


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", nargs="+", required=True)
    parser.add_argument("--labels", nargs="+", default=None)
    parser.add_argument("--repeats", type=int, default=5)
    parser.add_argument(
        "--blind", type=float, nargs="*", default=None,
        help="Known blind recovery per field, to report the search gap.",
    )
    parser.add_argument("--output", default="artifacts/vg_5k/anchor_bound.json")
    return parser.parse_args()


def standardise(matrix: np.ndarray) -> np.ndarray:
    values = matrix[~np.eye(len(matrix), dtype=bool)]
    out = (matrix - values.mean()) / values.std()
    np.fill_diagonal(out, 0.0)
    return out


def rows_against(field: np.ndarray, probe: np.ndarray, anchors: np.ndarray) -> np.ndarray:
    block = field[np.ix_(probe, anchors)]
    block = block - block.mean(1, keepdims=True)
    return block / block.std(1, keepdims=True).clip(1e-9)


def analyse(path: str, label: str, repeats: int, blind: float | None) -> dict:
    state = torch.load(path, map_location="cpu", weights_only=False)
    visual = standardise(state["visual_field"].double().numpy())
    text = standardise(state["text_field"].double().numpy())
    size = len(visual)
    mask = ~np.eye(size, dtype=bool)
    rho = float(np.corrcoef(visual[mask], text[mask])[0, 1])

    nearest, assigned = [], []
    for repeat in range(repeats):
        generator = np.random.default_rng(repeat)
        order = generator.permutation(size)
        half = size // 2
        anchors, probe = order[:half], order[half:]
        left = rows_against(visual, probe, anchors)
        right = rows_against(text, probe, anchors)
        similarity = left @ right.T / left.shape[1]
        truth = np.arange(len(probe))
        nearest.append(float((similarity.argmax(1) == truth).mean()))
        _, columns = linear_sum_assignment(-similarity)
        assigned.append(float((columns == truth).mean()))

    row = {
        "label": label,
        "correlation": rho,
        "probe_scenes": size - size // 2,
        "chance": 1.0 / (size - size // 2),
        "nearest_neighbour": float(np.mean(nearest)),
        "anchor_bound": float(np.mean(assigned)),
    }
    if blind is not None:
        row["blind_recovery"] = blind
        row["search_gap"] = float(np.mean(assigned)) - blind
    return row


def main() -> None:
    args = parse_args()
    labels = args.labels or [path.split("/")[-1] for path in args.fields]
    blinds = args.blind if args.blind else [None] * len(args.fields)
    rows = []
    for path, label, blind in zip(args.fields, labels, blinds):
        row = analyse(path, label, args.repeats, blind)
        rows.append(row)
        tail = (
            f"  blind={row['blind_recovery']:.3f}  gap={row['search_gap']:+.3f}"
            if "blind_recovery" in row else ""
        )
        print(
            f"{row['label']:<26} rho={row['correlation']:.3f}  "
            f"bound={row['anchor_bound']:.3f}{tail}",
            flush=True,
        )
    write_json(args.output, {
        "protocol": (
            "Half the scenes are anchors with their correspondence given. The "
            "rest are described by their field rows against the anchors and "
            "matched by Hungarian assignment. Upper-bounds blind recovery, "
            "which must additionally discover the anchor correspondence."
        ),
        "rows": rows,
        "reading": (
            "A low bound means the representation does not identify scenes and "
            "no solver will help. A high bound with low blind recovery means "
            "the information is present and the search or the landscape is at "
            "fault."
        ),
    })


if __name__ == "__main__":
    main()