Kac-rotated fast projection / run_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  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()