1import json
  2import time
  3import numpy as np
  4
  5
  6def fourier_concentration(n, band_fraction=0.5, spatial_fraction=0.75):
  7    x = np.arange(n)
  8    F = np.exp(-2j * np.pi * np.outer(x, x) / n) / np.sqrt(n)
  9    freqs = np.fft.fftfreq(n)
 10    k = max(1, int(round(band_fraction * n / 2)))
 11    band = np.abs(freqs) <= (k / n)
 12    Q = F.conj().T @ np.diag(band.astype(float)) @ F
 13    m = max(1, int(round(spatial_fraction * n)))
 14    mask = np.zeros(n); mask[:m] = 1.0
 15    return (np.diag(mask) @ Q @ np.diag(mask)).real
 16
 17
 18def kron_factors(factors):
 19    out = factors[0]
 20    for f in factors[1:]: out = np.kron(out, f)
 21    return out
 22
 23
 24def schatten(a, p):
 25    s = np.linalg.svd(a, compute_uv=False)
 26    return float(np.sum(s ** p) ** (1.0 / p))
 27
 28
 29def sequential_apply(x, factors, shape):
 30    y = x.reshape(shape)
 31    for axis, s in enumerate(factors):
 32        y = np.moveaxis(y, axis, 0)
 33        old = y.shape
 34        y = (s @ y.reshape(old[0], -1)).reshape(old)
 35        y = np.moveaxis(y, 0, axis)
 36    return y.reshape(-1)
 37
 38
 39def timed(fn, repeats=7):
 40    vals = []
 41    for _ in range(repeats):
 42        t = time.perf_counter(); fn(); vals.append(time.perf_counter() - t)
 43    return float(np.median(vals))
 44
 45
 46def main():
 47    rng = np.random.default_rng(1441)
 48    contraction_rows = []
 49    for d in [1, 2, 3, 4, 5]:
 50        factors = [fourier_concentration(9 + j, .55, .7) for j in range(d)]
 51        eigs = [np.linalg.eigvalsh(s) for s in factors]
 52        # Tensor eigenvalues are products: extrema follow directly since all are nonnegative.
 53        contraction_rows.append({
 54            "d": d, "factor_min_eigenvalue": float(min(v.min() for v in eigs)),
 55            "factor_max_eigenvalue": float(max(v.max() for v in eigs)),
 56            "global_min_eigenvalue": float(np.prod([v.min() for v in eigs])),
 57            "global_max_eigenvalue": float(np.prod([v.max() for v in eigs])),
 58            "predicted_max_bound": 1.0,
 59        })
 60
 61    tensor_rows = []
 62    for d in [2, 3, 4]:
 63        sizes = [6] * d
 64        factors = [fourier_concentration(n, .5, .67) for n in sizes]
 65        # A small explicit Kronecker matrix verifies application; no large SVD is needed.
 66        K = kron_factors(factors); x = rng.normal(size=6 ** d)
 67        exact = K @ x; sep = sequential_apply(x, factors, sizes)
 68        row = {"d": d, "apply_relative_error": float(np.linalg.norm(exact-sep)/np.linalg.norm(exact))}
 69        for p in [.5, 1., 2.]:
 70            lhs = schatten(K, p)
 71            rhs = float(np.prod([schatten(s, p) for s in factors]))
 72            row[f"schatten_p{p}_relative_error"] = abs(lhs-rhs) / max(rhs, 1e-15)
 73        tensor_rows.append(row)
 74
 75    scaling_rows = []
 76    for d in [2, 3, 4]:
 77        n = 8 if d <= 3 else 5
 78        factors = [fourier_concentration(n, .5, .75) for _ in range(d)]
 79        shape = [n] * d; N = n ** d
 80        K = kron_factors(factors); x = rng.normal(size=N)
 81        dense_t = timed(lambda: K @ x)
 82        axial_t = timed(lambda: sequential_apply(x, factors, shape))
 83        dense_ops = N * N; axial_ops = N * sum(shape)
 84        scaling_rows.append({"d": d, "tokens": N, "dense_parameters": N*N,
 85            "factor_parameters": d*n*n, "dense_matvec_ops": dense_ops,
 86            "axis_matvec_ops": axial_ops, "predicted_cost_ratio_dense_over_axis": dense_ops/axial_ops,
 87            "median_dense_seconds": dense_t, "median_axis_seconds": axial_t,
 88            "observed_time_ratio_dense_over_axis": dense_t/max(axial_t, 1e-15)})
 89
 90    n1 = n2 = 12; factors = [fourier_concentration(n1, .5, .75), fourier_concentration(n2, .5, .75)]
 91    K = kron_factors(factors); x = rng.normal(size=n1*n2)
 92    # Singular values of a tensor product are pairwise products, so this is exact and cheap.
 93    practical = {"tokens": n1*n2, "dense_parameters": int(K.size),
 94        "tensor_parameters": int(sum(s.size for s in factors)),
 95        "dense_apply_seconds": timed(lambda: K @ x),
 96        "tensor_apply_seconds": timed(lambda: sequential_apply(x, factors, [n1,n2])),
 97        "dense_largest_singular_value": float(np.prod([np.linalg.svd(s, compute_uv=False)[0] for s in factors])),
 98        "tensor_largest_singular_value": float(np.prod([np.linalg.svd(s, compute_uv=False)[0] for s in factors]))}
 99    print(json.dumps({"seed": 1441, "contraction_sweep": contraction_rows,
100        "tensor_identity_sweep": tensor_rows, "scaling_sweep": scaling_rows,
101        "practical_comparison": practical}, indent=2))
102
103if __name__ == "__main__": main()