summaryrefslogtreecommitdiff
path: root/worldalign/field_anatomy.py
blob: 9e04ec66bfedb9bbd0ce59b05e0f6f2563787807 (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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
"""What is the field correlation made of, and how much of it can matching use?

A puzzle motivates this. At 256 scenes the information-theoretic threshold is
0.29 and the fields correlate at 0.656, so the pairing is identifiable by a
wide margin, yet even at 64 scenes -- where the threshold is 0.51 and the
search space is trivially small -- recovery reaches 14%. Either the theory is
the wrong reference or the statistic is not measuring what the theory means.

The suspect is nuisance structure shared by both fields. Scene similarity has
a component that says only "this scene is busy, so it resembles everything",
and both modalities see it. That component inflates the correlation over all
pairs while carrying no information about which scene is which, because it is
the same for every candidate pairing of a scene to a partner of similar
busyness. Correlated-matching theory assumes the correlation is in the
exchangeable noise, so a nuisance-inflated statistic overstates what the
theory would predict.

This decomposes each field into a degree component -- the rank-one part
predicted by row and column means -- and the residual, and reports the
correlation of each. The residual correlation is the honest input to the
threshold, and if it falls far below 0.656 the recovery gap is explained and
the target moves.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", default="artifacts/vg_5k/natural_pipeline.pt")
    parser.add_argument("--label", default="natural")
    parser.add_argument("--output", default="artifacts/vg_5k/field_anatomy.json")
    return parser.parse_args()


def offdiagonal(matrix: np.ndarray) -> np.ndarray:
    return matrix[~np.eye(len(matrix), dtype=bool)]


def correlate(first: np.ndarray, second: np.ndarray) -> float:
    return float(np.corrcoef(offdiagonal(first), offdiagonal(second))[0, 1])


def degree_part(
    matrix: np.ndarray, sweeps: int = 200
) -> tuple[np.ndarray, np.ndarray]:
    """Split into the additive row/column-mean model and its residual.

    The two-way additive fit m + r_i + c_j is what a pure busyness effect
    produces: every entry explained by how prominent each of its two scenes
    is, with nothing said about the pair.

    The fit is swept to convergence rather than taken in one pass. With the
    diagonal excluded the design is unbalanced, so a single pass of row and
    column means leaves an O(1/n) piece of a pure degree effect behind --
    small, but it is exactly the quantity this function exists to remove.
    Alternating the two sweeps converges to the correct fit.
    """
    size = len(matrix)
    mask = ~np.eye(size, dtype=bool)
    work = np.where(mask, matrix.astype(np.float64), np.nan)
    grand = np.nanmean(work)
    rows = np.zeros((size, 1))
    cols = np.zeros((1, size))
    for _ in range(sweeps):
        residual = work - (grand + rows + cols)
        step = np.nanmean(residual, axis=1, keepdims=True)
        rows = rows + step
        residual = work - (grand + rows + cols)
        step_columns = np.nanmean(residual, axis=0, keepdims=True)
        cols = cols + step_columns
        if max(np.abs(step).max(), np.abs(step_columns).max()) < 1e-12:
            break
    fitted = grand + rows + cols
    return np.where(mask, fitted, 0.0), np.where(mask, work - fitted, 0.0)


def spectral_strip(matrix: np.ndarray, components: int) -> np.ndarray:
    """Remove the leading eigen-components, a stronger nuisance model."""
    symmetric = (matrix + matrix.T) / 2.0
    values, vectors = np.linalg.eigh(symmetric)
    order = np.argsort(np.abs(values))[::-1]
    keep = order[components:]
    return (vectors[:, keep] * values[keep]) @ vectors[:, keep].T


def main() -> None:
    args = parse_args()
    state = torch.load(args.fields, map_location="cpu", weights_only=False)
    visual = state["visual_field"].double().numpy()
    text = state["text_field"].double().numpy()
    size = len(visual)

    total = correlate(visual, text)
    visual_fit, visual_residual = degree_part(visual)
    text_fit, text_residual = degree_part(text)

    rows = [
        {"component": "raw field", "correlation": total},
        {"component": "degree model only", "correlation": correlate(visual_fit, text_fit)},
        {
            "component": "residual after degree",
            "correlation": correlate(visual_residual, text_residual),
        },
    ]
    for components in (1, 2, 5, 10):
        rows.append(
            {
                "component": f"residual after top-{components} eigen",
                "correlation": correlate(
                    spectral_strip(visual, components), spectral_strip(text, components)
                ),
            }
        )

    # how much of each field's own variance the nuisance model absorbs
    share = {
        "visual_degree_variance_share": float(
            offdiagonal(visual_fit).var() / offdiagonal(visual).var()
        ),
        "text_degree_variance_share": float(
            offdiagonal(text_fit).var() / offdiagonal(text).var()
        ),
    }

    residual_rho = rows[2]["correlation"]
    threshold = float(np.sqrt(4 * np.log(size) / size))
    summary = {
        "protocol": (
            "Both fields split into an additive row/column-mean model and its "
            "residual, then correlated component by component. Hidden pairs "
            "are read only to align the two fields."
        ),
        "label": args.label,
        "size": size,
        "rows": rows,
        **share,
        "it_threshold_at_size": threshold,
        "residual_above_it_threshold": bool(residual_rho > threshold),
        "reading": (
            "If the residual correlation is far below the raw correlation, the "
            "headline statistic is inflated by nuisance structure that carries "
            "no matching information, and the residual is the number that "
            "should be compared against the recovery threshold."
        ),
    }
    for row in rows:
        print(f"{row['component']:<34} rho={row['correlation']:.4f}")
    print(json.dumps(share))
    print(f"IT threshold at N={size}: {threshold:.4f}")
    write_json(args.output, summary)


if __name__ == "__main__":
    main()