1import json
  2import math
  3import os
  4import numpy as np
  5from scipy.linalg import hadamard
  6
  7SEED = 2532
  8rng = np.random.default_rng(SEED)
  9
 10
 11def signed_perm(d, rng):
 12    p = rng.permutation(d)
 13    s = rng.choice(np.array([-1.0, 1.0]), size=d)
 14    # Sigma z = s * z[p].
 15    S = np.zeros((d, d))
 16    S[np.arange(d), p] = s
 17    return S
 18
 19
 20def make_u(d, rng):
 21    H = hadamard(d).astype(float) / math.sqrt(d)
 22    A = H
 23    B = H
 24    return A @ signed_perm(d, rng) @ B
 25
 26
 27def coherence_sweep(dims=(16, 32, 64, 128, 256), trials=200):
 28    rows = []
 29    for d in dims:
 30        vals = []
 31        ortho = []
 32        for _ in range(trials):
 33            U = make_u(d, rng)
 34            vals.append(math.sqrt(d) * np.max(np.abs(U)))
 35            ortho.append(np.linalg.norm(U.T @ U - np.eye(d), ord='fro'))
 36        vals = np.asarray(vals)
 37        predicted = math.sqrt(2.0 * math.log(d))
 38        rows.append({
 39            'd': d,
 40            'coherence_mu_mean': float(vals.mean()),
 41            'coherence_mu_p95': float(np.quantile(vals, .95)),
 42            'predicted_sqrt_2logd': predicted,
 43            'ratio_mean_to_prediction': float(vals.mean() / predicted),
 44            'orthogonality_fro_mean': float(np.mean(ortho)),
 45        })
 46    return rows
 47
 48
 49def quantize(z, bits=4):
 50    q = 2 ** (bits - 1) - 1
 51    delta = np.max(np.abs(z), axis=-1, keepdims=True) / max(q, 1)
 52    delta = np.maximum(delta, 1e-12)
 53    return np.clip(np.rint(z / delta), -q, q) * delta
 54
 55
 56def make_vectors(d, n, sparsity, rng):
 57    # Each vector has k active coordinates with Gaussian values and one controlled
 58    # outlier scale. This stresses block-max quantization while retaining random orientation.
 59    k = max(1, int(round(d * sparsity)))
 60    X = np.zeros((n, d))
 61    for i in range(n):
 62        idx = rng.choice(d, size=k, replace=False)
 63        X[i, idx] = rng.normal(size=k)
 64    # Normalize every block so comparisons have identical input energy.
 65    X /= np.linalg.norm(X, axis=1, keepdims=True)
 66    return X
 67
 68
 69def quantization_sweep(dims=(64, 128, 256), sparsities=(1/256, 1/64, 1/16, 1/4, 1.0), n=1000, bits=4):
 70    rows = []
 71    for d in dims:
 72        # one transform per dimension, frozen as proposed
 73        U = make_u(d, rng)
 74        for sp in sparsities:
 75            X = make_vectors(d, n, sp, rng)
 76            # Baseline: coordinate quantization; permutation-only is mathematically
 77            # identical to this max-scale quantizer and is included as a check.
 78            Y0 = quantize(X, bits)
 79            perm = rng.permutation(d)
 80            Xperm = X[:, perm]
 81            Yp = quantize(Xperm, bits)[:, np.argsort(perm)]
 82            # Full preconditioner and exact inverse.
 83            Z = X @ U.T
 84            Zq = quantize(Z, bits)
 85            Xhat = Zq @ U
 86            mse0 = float(np.mean((Y0 - X) ** 2))
 87            msep = float(np.mean((Yp - X[:, rng.permutation(d)]) ** 2))
 88            msem = float(np.mean((Xhat - X) ** 2))
 89            # normalized by signal energy, which is 1 per row
 90            rows.append({
 91                'd': d, 'sparsity': sp, 'baseline_nmse': mse0,
 92                'perm_nmse': msep, 'mixed_nmse': msem,
 93                'mixed_gain_vs_baseline': mse0 / max(msem, 1e-30),
 94                'baseline_max_abs_mean': float(np.mean(np.max(np.abs(X), axis=1))),
 95                'mixed_max_abs_mean': float(np.mean(np.max(np.abs(Z), axis=1))),
 96            })
 97    return rows
 98
 99
100def sampling_check(d=128, fractions=(1/32, 1/8, 1/2), trials=3000):
101    # For unbiased coordinate sampling, theory predicts E||xhat-x||^2 = d/m - 1
102    # for every unit vector. Mixing should additionally reduce the largest sampled
103    # coordinate energy for sparse inputs, as predicted by incoherence.
104    U = make_u(d, rng)
105    x = np.zeros(d); x[0] = 1.0
106    z = U @ x
107    rows = []
108    for frac in fractions:
109        m = max(1, int(round(d * frac)))
110        errs0, errsm = [], []
111        for _ in range(trials):
112            I = rng.choice(d, size=m, replace=False)
113            y0 = np.zeros(d); y0[I] = x[I] * d / m
114            yz = np.zeros(d); yz[I] = z[I] * d / m
115            xr = U.T @ yz
116            errs0.append(np.sum((y0-x)**2))
117            errsm.append(np.sum((xr-x)**2))
118        rows.append({'d': d, 'm': m, 'predicted_nmse': d/m-1,
119                     'identity_mean_nmse': float(np.mean(errs0)),
120                     'mixed_mean_nmse': float(np.mean(errsm)),
121                     'mixed_p95_nmse': float(np.quantile(errsm,.95)),
122                     'identity_max_coordinate_energy': 1.0,
123                     'mixed_max_coordinate_energy': float(np.max(z*z))})
124    return rows
125
126
127def main():
128    coh = coherence_sweep()
129    quant = quantization_sweep()
130    # A simple quantitative prediction from incoherence: for a k-sparse unit vector,
131    # a dense incoherent transform should have max coordinate O(sqrt(log d/d)),
132    # while identity has max=1. Report the observed ratio at the sparsest setting.
133    prediction = []
134    for d in (64, 128, 256):
135        r = next(x for x in quant if x['d'] == d and abs(x['sparsity'] - 1/256) < 1e-12)
136        predicted_max = math.sqrt(2 * math.log(d) / d)
137        prediction.append({
138            'd': d,
139            'predicted_mixed_max_scale': predicted_max,
140            'observed_mixed_max_abs': r['mixed_max_abs_mean'],
141            'observed_identity_max_abs': r['baseline_max_abs_mean'],
142            'predicted_identity_max_abs': 1.0,
143            'observed_mixed_to_identity': r['mixed_max_abs_mean'] / r['baseline_max_abs_mean'],
144        })
145    # Prediction 1: coherence is approximately sqrt(2 log d / d), up to a
146    # slowly varying constant; estimate the observed log-log slope after removing log.
147    ds = np.array([x['d'] for x in coh], float)
148    ys = np.array([x['coherence_mu_mean'] for x in coh]) / np.sqrt(ds)
149    slope = float(np.polyfit(np.log(ds), np.log(ys), 1)[0])
150    sampling = sampling_check()
151    result = {'seed': SEED, 'coherence': coh, 'coherence_scaled_loglog_slope': slope,
152              'quantization': quant, 'sparse_spreading_prediction': prediction,
153              'sampling': sampling}
154    with open('results.json', 'w') as f:
155        json.dump(result, f, indent=2)
156    print(json.dumps(result, indent=2))
157
158
159if __name__ == '__main__':
160    main()