1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6
  7def concentration_matrix(T, W):
  8    i = np.arange(T)
  9    d = i[:, None] - i[None, :]
 10    C = np.empty((T, T), dtype=float)
 11    nz = d != 0
 12    C[nz] = np.sin(2.0 * np.pi * W * d[nz]) / (np.pi * d[nz])
 13    C[~nz] = 2.0 * W
 14    return (C + C.T) * 0.5
 15
 16
 17def fourier_basis(T, r):
 18    # Complex unitary DFT columns, ordered by distance from zero frequency.
 19    n = np.arange(T)
 20    freqs = np.fft.fftfreq(T)
 21    order = np.argsort(np.abs(freqs), kind="stable")[:r]
 22    U = np.exp(2j * np.pi * np.outer(n, freqs[order])) / np.sqrt(T)
 23    return U
 24
 25
 26def projection_error(X, U):
 27    # X: samples x T; U: T x r orthonormal (possibly complex)
 28    Z = X @ U
 29    Xhat = Z @ U.conj().T
 30    return float(np.mean(np.sum(np.abs(X - Xhat) ** 2, axis=1) /
 31                     np.maximum(np.sum(np.abs(X) ** 2, axis=1), 1e-12)))
 32
 33
 34def main():
 35    rng = np.random.default_rng(2368)
 36    T, W = 96, 0.22
 37    delta = 0.1
 38    C = concentration_matrix(T, W)
 39    evals, Uall = np.linalg.eigh(C)
 40    order = np.argsort(evals)[::-1]
 41    evals, Uall = evals[order], Uall[:, order]
 42
 43    # The finite concentration operator should be PSD, have eigenvalues in [0,1],
 44    # and have effective dimension approximately 2WT.
 45    c = 2.0 * W * T
 46    barL = math.log((1.0 - delta) / delta)
 47    asym_rank = c + barL / math.pi**2 * math.log(max(4.0 * math.pi**2 * c / barL, 1.000001))
 48    r_emp = int(np.sum(evals > delta))
 49    r_asym = int(np.clip(math.ceil(asym_rank), 1, T))
 50
 51    # Draw finite-window samples from the band-limited covariance C. In this
 52    # distribution, the DPSS eigenvectors are the optimal rank-r coordinates.
 53    # A tiny jitter makes the square root robust to roundoff only.
 54    pos = np.maximum(evals, 0.0)
 55    X = rng.standard_normal((1200, T)) @ (Uall * np.sqrt(pos)).T
 56
 57    ranks = [8, 12, 16, 20, 24, 28, 32, 40, 48]
 58    rows = []
 59    for r in ranks:
 60        Ud = Uall[:, :r]
 61        Uf = fourier_basis(T, r)
 62        Ur = np.linalg.qr(rng.standard_normal((T, r)))[0]
 63        rows.append({
 64            "rank": r,
 65            "dpss_error": projection_error(X, Ud),
 66            "fourier_error": projection_error(X.astype(complex), Uf),
 67            "random_error": projection_error(X, Ur),
 68        })
 69
 70    # At the theorem-inspired threshold, also report the retained energy and
 71    # the predicted attention quadratic-work reduction.
 72    r = r_asym
 73    retained = float(np.sum(evals[:r]) / np.sum(evals))
 74    threshold_retained = float(np.sum(evals[:r_emp]) / np.sum(evals))
 75    result = {
 76        "config": {"T": T, "W": W, "delta": delta, "samples": len(X)},
 77        "math_check": {
 78            "min_eigenvalue": float(evals[-1]),
 79            "max_eigenvalue": float(evals[0]),
 80            "eigenvalues_in_unit_interval": bool(evals[0] <= 1.0 + 1e-10 and evals[-1] >= -1e-10),
 81            "trace": float(np.trace(C)),
 82            "time_bandwidth_2WT": c,
 83            "asymptotic_rank": float(asym_rank),
 84            "asymptotic_rank_clipped_ceiling": r_asym,
 85            "empirical_rank_above_delta": r_emp,
 86            "empirical_rank_relative_error": float(abs(r_emp - c) / c),
 87            "energy_retained_at_asym_rank": retained,
 88            "energy_retained_above_delta_rank": threshold_retained,
 89        },
 90        "curves": rows,
 91        "comparison_at_asym_rank": {
 92            "rank": r,
 93            "dpss_error": rows[[q["rank"] for q in rows].index(r)]["dpss_error"] if r in ranks else projection_error(X, Uall[:, :r]),
 94            "fourier_error": rows[[q["rank"] for q in rows].index(r)]["fourier_error"] if r in ranks else projection_error(X.astype(complex), fourier_basis(T, r)),
 95            "random_error": rows[[q["rank"] for q in rows].index(r)]["random_error"] if r in ranks else None,
 96            "quadratic_attention_work_fraction": float((r / T) ** 2),
 97            "quadratic_attention_work_reduction": float(1.0 - (r / T) ** 2),
 98        },
 99    }
100    Path("results.json").write_text(json.dumps(result, indent=2))
101    print(json.dumps(result, indent=2))
102
103
104if __name__ == "__main__":
105    main()