summaryrefslogtreecommitdiff
path: root/worldalign/rho_at_full_width.py
blob: c30c5cc6e0bce23af80d977b95aae6cc54d2ee1f (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
"""The other axis: recovery against correlation with the shared width intact.

The truncation ladder varied the width at nearly fixed correlation. This
varies the correlation at full width, by adding independent noise to both
fields, and the two together give the picture the single statistic could not.

The control matters because it guards against overcorrecting. Showing that
width is a separate necessary condition does not show that correlation stopped
mattering, and natural data is short on both -- correlation 0.725 against 0.93,
width 15 against 26. If full-width fields still recover at a correlation near
0.72, then width is what natural data lacks; if they do not, both gaps are real
and the correlation remains a target rather than a discredited one.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json
from .shared_rank import spectrum
from .spectral_match import grampa
from .synth_fast_gate import ClosedFormEnergy, all_swaps, steepest_descent
from .synth_triangle_gate import standardized


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", default="artifacts/synth_v1/fields_tier0_ws_256.pt")
    parser.add_argument("--trials", type=int, default=3)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument(
        "--noise", type=float, nargs="+",
        default=[0.0, 0.2, 0.35, 0.5, 0.7, 0.9, 1.2],
    )
    parser.add_argument("--output", default="artifacts/synth_v1/rho_at_full_width.json")
    return parser.parse_args()


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


def normalise(matrix: np.ndarray) -> np.ndarray:
    values = offdiagonal(matrix)
    out = (matrix - values.mean()) / values.std()
    np.fill_diagonal(out, 0.0)
    return out


def symmetric_noise(size: int, generator) -> np.ndarray:
    raw = generator.normal(size=(size, size))
    out = (raw + raw.T) / np.sqrt(2.0)
    np.fill_diagonal(out, 0.0)
    return out


def shared_count(visual: np.ndarray, text: np.ndarray, width: int = 32) -> int:
    _, visual_vectors = spectrum(visual)
    _, text_vectors = spectrum(text)
    cosines = np.clip(
        np.linalg.svd(
            visual_vectors[:, :width].T @ text_vectors[:, :width], compute_uv=False
        ), 0.0, 1.0,
    )
    return int((cosines > 0.7).sum())


def main() -> None:
    args = parse_args()
    state = torch.load(args.fields, map_location="cpu", weights_only=False)
    visual_clean = normalise(state["visual_field"].double().numpy())
    text_clean = normalise(state["text_field"].double().numpy())
    size = len(visual_clean)
    device = torch.device(args.device)
    swaps = all_swaps(size, device)

    rows = []
    for level in args.noise:
        correlations, accuracies, widths = [], [], []
        for trial in range(args.trials):
            generator = np.random.default_rng(1000 + trial)
            visual = normalise(
                visual_clean + level * symmetric_noise(size, generator)
            )
            text = normalise(text_clean + level * symmetric_noise(size, generator))
            mask = ~np.eye(size, dtype=bool)
            correlations.append(float(np.corrcoef(visual[mask], text[mask])[0, 1]))
            widths.append(shared_count(visual, text))

            hidden = generator.permutation(size)
            shuffled = text[np.ix_(hidden, hidden)]
            columns = grampa(visual, shuffled, 1.0)
            energy = ClosedFormEnergy(
                standardized(torch.from_numpy(shuffled).to(device)).float(),
                standardized(torch.from_numpy(visual).to(device)).float(),
                1.0, 1.0, 256,
            )
            final, _ = steepest_descent(
                energy, torch.from_numpy(columns.copy()).to(device), swaps, 4000
            )
            accuracies.append(
                float((hidden[final.cpu().numpy()] == np.arange(size)).mean())
            )
        row = {
            "noise": level,
            "correlation": float(np.mean(correlations)),
            "shared_directions": float(np.mean(widths)),
            "recovery": float(np.mean(accuracies)),
        }
        rows.append(row)
        print(
            f"noise={level:<5} rho={row['correlation']:.3f}  "
            f"shared={row['shared_directions']:5.1f}  "
            f"recovery={row['recovery']:.3f}",
            flush=True,
        )

    summary = {
        "protocol": (
            "Independent symmetric noise added to both fields, so the shared "
            "spectrum stays wide while the correlation falls. Hidden pairing "
            "generated per trial and read only for scoring."
        ),
        "size": size,
        "chance": 1.0 / size,
        "rows": rows,
    }
    print(json.dumps({"chance": 1.0 / size}))
    write_json(args.output, summary)


if __name__ == "__main__":
    main()