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
|
"""World matching: spectral initialisation plus energy refinement.
The two solvers are complementary. A spectral solver reads the coarse
correspondence out of the eigenstructure of two relation fields in
polynomial time and without any search, but its rounding is noisy. Exact
steepest descent on the closed-form energy repairs the rounding but
cannot find the basin from a random start. Composed, they recover the
hidden assignment.
Neither stage sees a pair. The relation fields are built independently
per modality; the hidden permutation is applied to the text field and
read back only to score.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from .common import seed_everything, write_json
from .spectral_match import grampa, spectral_profile, umeyama
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", required=True)
parser.add_argument("--eta", type=float, default=1.0)
parser.add_argument("--pair-weight", type=float, default=1.0)
parser.add_argument("--triangle-weight", type=float, default=1.0)
parser.add_argument("--refine-steps", type=int, default=4000)
parser.add_argument("--chunk", type=int, default=256)
parser.add_argument("--restarts", type=int, default=1)
parser.add_argument("--device", default="cuda:3")
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--output", default="")
return parser.parse_args()
def normalise(matrix: np.ndarray) -> np.ndarray:
size = len(matrix)
mask = ~np.eye(size, dtype=bool)
values = matrix[mask]
out = (matrix - values.mean()) / values.std()
np.fill_diagonal(out, 0.0)
return out
def main() -> None:
args = parse_args()
seed_everything(args.seed)
state = torch.load(args.fields, map_location="cpu", weights_only=False)
visual = normalise(state["visual_field"].double().numpy().copy())
text = normalise(state["text_field"].double().numpy().copy())
size = len(visual)
generator = np.random.default_rng(args.seed)
hidden = generator.permutation(size)
shuffled = text[np.ix_(hidden, hidden)]
truth = np.arange(size)
def accuracy(assignment: np.ndarray) -> float:
return float((hidden[assignment] == truth).mean())
device = torch.device(args.device)
text_tensor = standardized(torch.from_numpy(shuffled).to(device)).float()
visual_tensor = standardized(torch.from_numpy(visual).to(device)).float()
energy = ClosedFormEnergy(
text_tensor,
visual_tensor,
args.pair_weight,
args.triangle_weight,
args.chunk,
)
swaps = all_swaps(size, device)
truth_permutation = torch.from_numpy(np.argsort(hidden)).to(device)
true_energy = float(energy.energy(truth_permutation[None])[0])
report = {
"protocol": (
"Fields are built per modality without pairs; the hidden "
"permutation is applied to the text field and used only to "
"score. Spectral initialisation is followed by exact "
"steepest descent on the closed-form energy."
),
"samples": size,
"field_correlation_at_truth": float(
np.corrcoef(
visual[~np.eye(size, dtype=bool)],
text[~np.eye(size, dtype=bool)],
)[0, 1]
),
"true_energy": true_energy,
"chance": 1.0 / size,
"spectra": {
"visual": spectral_profile(visual),
"text": spectral_profile(shuffled),
},
"stages": {},
}
spectral = grampa(visual, shuffled, args.eta)
report["stages"]["spectral"] = {"accuracy": accuracy(spectral)}
initial = torch.from_numpy(spectral.copy()).to(device)
refined, value = steepest_descent(energy, initial, swaps, args.refine_steps)
report["stages"]["spectral_plus_refinement"] = {
"accuracy": accuracy(refined.cpu().numpy()),
"energy": value,
"energy_over_true": value / abs(true_energy) - true_energy / abs(true_energy),
}
quenches = []
for restart in range(args.restarts):
start = torch.from_numpy(generator.permutation(size)).to(device)
final, final_value = steepest_descent(
energy, start, swaps, args.refine_steps
)
quenches.append(
{
"accuracy": accuracy(final.cpu().numpy()),
"energy": final_value,
}
)
report["stages"]["random_start_refinement"] = quenches
print(json.dumps({key: value for key, value in report["stages"].items()}))
if args.output:
write_json(args.output, report)
print(f"Wrote {args.output}")
if __name__ == "__main__":
main()
|
