Kac-rotated fast projection / run_experiment.py
Failed on benchmark
1import json, math, time
2from pathlib import Path
3import numpy as np
4
5SEED = 2454
6rng = np.random.default_rng(SEED)
7
8
9def make_rotations(n, T, seed):
10 r = np.random.default_rng(seed)
11 ij = r.integers(0, n, size=(T, 2))
12 # Resample equal pairs; sorting gives valid i<j.
13 eq = ij[:, 0] == ij[:, 1]
14 while np.any(eq):
15 ij[eq, 1] = r.integers(0, n, size=int(eq.sum()))
16 eq = ij[:, 0] == ij[:, 1]
17 a = np.minimum(ij[:, 0], ij[:, 1])
18 b = np.maximum(ij[:, 0], ij[:, 1])
19 th = r.uniform(0.0, 2.0 * np.pi, size=T)
20 return a.astype(np.int32), b.astype(np.int32), np.cos(th), np.sin(th)
21
22
23def kac_apply(x, rotations):
24 # x has shape (batch,n) or (n,); this is the stated streamed update.
25 a, b, c, s = rotations
26 z = np.array(x, dtype=np.float64, copy=True)
27 for i, j, cc, ss in zip(a, b, c, s):
28 u = z[..., i].copy()
29 v = z[..., j].copy()
30 z[..., i] = cc * u + ss * v
31 z[..., j] = -ss * u + cc * v
32 return z
33
34
35def haar_projection(n, m, seed):
36 r = np.random.default_rng(seed)
37 g = r.normal(size=(n, n))
38 q, rr = np.linalg.qr(g)
39 # Make the QR convention deterministic up to column signs; either is Haar.
40 q *= np.sign(np.diag(rr))[None, :]
41 return np.sqrt(n / m) * q[:m, :]
42
43
44def gaussian_projection(n, m, seed):
45 return np.random.default_rng(seed).normal(size=(m, n)) / math.sqrt(m)
46
47
48def metrics(values):
49 e = values - 1.0
50 return {
51 "mean": float(values.mean()),
52 "mean_abs_error": float(np.mean(np.abs(e))),
53 "std": float(values.std()),
54 "p95_abs_error": float(np.quantile(np.abs(e), .95)),
55 "mse": float(np.mean(e * e)),
56 }
57
58
59def main():
60 n = 64
61 m = 16
62 x = rng.normal(size=n)
63 x /= np.linalg.norm(x)
64 # Haar benchmark uses independent Haar directions, with the same fixed x norm.
65 trials = 500
66 haar_vals = np.empty(trials)
67 for k in range(trials):
68 q = haar_projection(n, m, 10000 + k)
69 haar_vals[k] = np.sum((q @ x) ** 2)
70 haar = metrics(haar_vals)
71
72 # Prediction 1: every finite Kac product is orthogonal, independently of T.
73 orth = []
74 # Prediction 2: mean scaled subset energy is approximately one and approaches one.
75 # Prediction 3: quadratic-statistic spread approaches the Haar spread as T grows.
76 sweep = [0, n // 4, n // 2, n, 2*n, 4*n, 8*n, 16*n]
77 rows = []
78 for T in sweep:
79 vals = np.empty(trials)
80 defects = np.empty(30)
81 for k in range(trials):
82 rot = make_rotations(n, T, 20000 + 1000 * k + T)
83 z = kac_apply(x, rot)
84 vals[k] = (n / m) * np.sum(z[:m] ** 2)
85 if k < len(defects):
86 # Applying the same product to a few basis vectors estimates Q^T Q.
87 Q = np.column_stack([kac_apply(np.eye(n)[j], rot) for j in range(n)])
88 defects[k] = np.max(np.abs(Q.T @ Q - np.eye(n)))
89 mm = metrics(vals)
90 pred_mean = 1.0 + (float((n/m) * np.sum(x[:m]**2)) - 1.0) * ((n-1.0)/n)**T
91 rows.append({"T": T, "predicted_mean": pred_mean, "mean_prediction_abs_error": abs(mm["mean"] - pred_mean), **mm, "std_over_haar_std": float(mm["std"] / haar["std"]),
92 "orthogonality_max_abs": float(defects.max())})
93 orth.append(float(defects.max()))
94
95 # Distortion benchmark on a common batch of pair differences.
96 batch = 128
97 X = rng.normal(size=(batch, n))
98 X /= np.linalg.norm(X, axis=1, keepdims=True)
99 pairs = X[::2] - X[1::2]
100 # Compare the same fixed Kac transform to dense references at practical T.
101 Tbench = 8 * n
102 rot = make_rotations(n, Tbench, 777)
103 t0 = time.perf_counter(); Yk = np.sqrt(n/m) * kac_apply(X, rot)[:, :m]; kac_time = time.perf_counter()-t0
104 G = gaussian_projection(n, m, 778)
105 t0 = time.perf_counter(); Yg = X @ G.T; gauss_time = time.perf_counter()-t0
106 H = haar_projection(n, m, 779)
107 t0 = time.perf_counter(); Yh = X @ H.T; haar_time = time.perf_counter()-t0
108 def pair_dist(Y):
109 return np.sum((Y[::2]-Y[1::2])**2, axis=1) / np.sum(pairs*pairs, axis=1)
110 distortion = {"Kac": metrics(pair_dist(Yk)), "Gaussian": metrics(pair_dist(Yg)), "Haar": metrics(pair_dist(Yh))}
111 # Repeat timings enough to avoid timer noise, while keeping the literal stream honest.
112 reps = 5
113 t0=time.perf_counter()
114 for _ in range(reps): kac_apply(X, rot)
115 kac_time=(time.perf_counter()-t0)/reps
116 t0=time.perf_counter()
117 for _ in range(reps): X @ G.T
118 gauss_time=(time.perf_counter()-t0)/reps
119 explicit_triple_bytes = int(Tbench * (4 + 4 + 8))
120 dense_orthogonal_bytes = int(n * n * 8)
121 storage = {"Kac_runtime_arrays_bytes": int(sum(a.nbytes for a in rot)),
122 "Kac_explicit_i_j_theta_bytes": explicit_triple_bytes,
123 "dense_projection_bytes": int(G.nbytes),
124 "dense_full_orthogonal_bytes": dense_orthogonal_bytes,
125 "dense_full_over_explicit_triples": float(dense_orthogonal_bytes / explicit_triple_bytes),
126 "dense_projection_over_runtime_arrays": float(G.nbytes / sum(a.nbytes for a in rot))}
127 result = {
128 "seed": SEED, "n": n, "m": m, "trials": trials,
129 "predictions": {
130 "P1": "orthogonality defect remains at floating-point roundoff for every T",
131 "P2": "scaled quadratic norm has mean near 1",
132 "P3": "quadratic-statistic spread moves toward the Haar reference as T increases"
133 },
134 "haar_reference": haar, "sweep": rows,
135 "prediction_checks": {
136 "P1_max_defect_all_T": float(max(orth)),
137 "P2_mean_range": [float(min(r["mean"] for r in rows)), float(max(r["mean"] for r in rows))],
138 "P2_max_abs_error_vs_analytic_mean": float(max(r["mean_prediction_abs_error"] for r in rows)),
139 "P3_std_ratio_T0": rows[0]["std_over_haar_std"],
140 "P3_std_ratio_T8n": next(r["std_over_haar_std"] for r in rows if r["T"] == 8*n),
141 "P3_std_ratio_T16n": rows[-1]["std_over_haar_std"]
142 },
143 "distortion": distortion,
144 "timing_seconds_per_batch": {"Kac_literal": kac_time, "Gaussian_dense": gauss_time, "Haar_dense_not_timed": haar_time},
145 "storage": storage
146 }
147 Path("results.json").write_text(json.dumps(result, indent=2))
148 print(json.dumps(result, indent=2))
149
150if __name__ == "__main__":
151 main()