1import json
 2import time
 3from pathlib import Path
 4import numpy as np
 5from scipy.optimize import linear_sum_assignment
 6
 7
 8def standardize_features(F):
 9    F = np.asarray(F, dtype=np.float64)
10    F = F - F.mean(axis=0, keepdims=True)
11    s = F.std(axis=0, keepdims=True)
12    s[s < 1e-12] = 1.0
13    return F / s
14
15
16def cosine_score(X, Y):
17    Xn = X / np.maximum(np.linalg.norm(X, axis=1, keepdims=True), 1e-12)
18    Yn = Y / np.maximum(np.linalg.norm(Y, axis=1, keepdims=True), 1e-12)
19    return Xn @ Yn.T
20
21
22def tree_features(X, max_depth=5, branches=(2, 3, 4), normalize=True):
23    """Finite rooted-tree contractions computed from the sample Gram matrix.
24
25    A coordinate vertex of degree two contracts to A=XX^T/d. Repeated A
26    multiplication gives rooted path contractions; powers of A1 represent
27    wide rooted trees with repeated branches. Thus all features are exactly
28    invariant under right multiplication X -> XQ for orthogonal Q.
29    """
30    X = np.asarray(X, dtype=np.float64)
31    A = X @ X.T / max(X.shape[1], 1)
32    one = np.ones((X.shape[0], 1), dtype=np.float64)
33    messages, v = [], one
34    for _ in range(max_depth):
35        v = A @ v
36        messages.append(v[:, 0])
37    root = (A @ one)[:, 0]
38    for b in branches:
39        messages.append(root ** b)
40    F = np.stack(messages, axis=1)
41    return standardize_features(F) if normalize else F
42
43
44def tree_score(X, Y):
45    FX, FY = tree_features(X), tree_features(Y)
46    return FX @ FY.T / FX.shape[1]
47
48
49def hungarian_accuracy(score, truth=None):
50    r, c = linear_sum_assignment(-score)
51    if truth is None:
52        truth = np.arange(score.shape[0])
53    return float(np.mean(c == np.asarray(truth)[r]))
54
55
56def random_orthogonal(d, rng):
57    q, r = np.linalg.qr(rng.normal(size=(d, d)))
58    q *= np.sign(np.diag(r))[None, :]
59    return q
60
61
62def make_pair(n, d, rho, rng):
63    X = rng.normal(size=(n, d))
64    Q = random_orthogonal(d, rng)
65    # Y has row-wise correlation rho with the rotated matched X.
66    Y = rho * (X @ Q) + np.sqrt(max(1 - rho*rho, 0.0)) * rng.normal(size=(n, d))
67    return X, Y
68
69
70def run(seed=7, n=96, d=64, rhos=(0.55, 0.65, 0.75, 0.85), repeats=8):
71    rng = np.random.default_rng(seed)
72    invariant_errors, rows = [], []
73    X = rng.normal(size=(n, d)); Q = random_orthogonal(d, rng)
74    f0 = tree_features(X, normalize=False)
75    f1 = tree_features(X @ Q, normalize=False)
76    invariant_errors.append(float(np.max(np.abs(f0-f1)) / max(np.max(np.abs(f0)), 1e-12)))
77    for rho in rhos:
78        ca, ta = [], []
79        for _ in range(repeats):
80            X, Y = make_pair(n, d, rho, rng)
81            t = time.perf_counter(); cs = cosine_score(X, Y); tc = time.perf_counter()-t
82            t = time.perf_counter(); ts = tree_score(X, Y); tt = time.perf_counter()-t
83            ca.append(hungarian_accuracy(cs)); ta.append(hungarian_accuracy(ts))
84        rows.append({'rho': rho, 'cosine_mean': float(np.mean(ca)), 'cosine_std': float(np.std(ca)),
85                     'tree_mean': float(np.mean(ta)), 'tree_std': float(np.std(ta)),
86                     'cosine_sec': tc, 'tree_sec': tt})
87    out = {'invariance_relative_max_error': invariant_errors[0], 'rows': rows}
88    Path('results.json').write_text(json.dumps(out, indent=2))
89    print(json.dumps(out, indent=2))
90    return out
91
92
93if __name__ == '__main__':
94    run()