1import json
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1109
  6rng = np.random.default_rng(SEED)
  7D = 256
  8A = 1.0
  9u = (np.arange(D) + 0.5) / D
 10omega = 2*np.pi * u**(1/(A+1))
 11q = np.ones(D) / D
 12
 13
 14def return_proxy(lam, weights, ks):
 15    lam = np.asarray(lam, complex)
 16    weights = np.asarray(weights, float)
 17    weights = weights / weights.sum()
 18    amp = np.sum(weights[None, :] * lam[None, :] ** ks[:, None], axis=1)
 19    return np.abs(amp)**2
 20
 21
 22def fit_slope(ks, vals, lo=20, hi=180):
 23    sel = (ks >= lo) & (ks <= hi) & (vals > 1e-14)
 24    if sel.sum() < 3:
 25        return float("nan")
 26    return float(np.polyfit(np.log(ks[sel]), np.log(vals[sel]), 1)[0])
 27
 28ks = np.arange(1, 401)
 29rho0 = 0.9998
 30teacher_lam = rho0 * np.exp(1j * omega)
 31R_teacher = return_proxy(teacher_lam, q, ks)
 32
 33
 34def select_frequency(M):
 35    return np.linspace(0, D - 1, M).round().astype(int)
 36
 37
 38def select_random(M):
 39    return np.sort(rng.choice(D, M, replace=False))
 40
 41
 42def select_magnitude(M):
 43    # Equal modal magnitudes make magnitude pruning unable to distinguish modes.
 44    return np.arange(M)
 45
 46
 47def relative_log_error(a, b):
 48    return float(np.sqrt(np.mean((np.log(a + 1e-16) - np.log(b + 1e-16)) ** 2)))
 49
 50# Compression comparison. Selected modes inherit teacher weights and are renormalized.
 51compression = []
 52for M in [8, 16, 32, 64]:
 53    for name, selector in [("frequency", select_frequency), ("random", select_random), ("magnitude", select_magnitude)]:
 54        idx = selector(M)
 55        R = return_proxy(teacher_lam[idx], q[idx], ks)
 56        compression.append({
 57            "M": M, "method": name,
 58            "log_rmse": relative_log_error(R_teacher, R),
 59            "slope": fit_slope(ks, R),
 60            "teacher_slope": fit_slope(ks, R_teacher),
 61            "slope_relative_error": abs(fit_slope(ks, R) - fit_slope(ks, R_teacher)) / max(abs(fit_slope(ks, R_teacher)), 1e-12),
 62        })
 63
 64# Prediction 1: stable radius gives an exponential envelope with time constant
 65# tau = -1/(2 log rho), because R(k) contains rho^(2k).
 66radii = [0.90, 0.95, 0.99, 0.999]
 67stability = []
 68for rho in radii:
 69    R = return_proxy(rho * np.exp(1j * omega[:32]), q[:32], ks)
 70    # Fit log envelope after averaging oscillations by a robust upper quantile in bins.
 71    bins = np.array_split(np.arange(len(ks)), 20)
 72    kb = np.array([np.mean(ks[b]) for b in bins])
 73    rb = np.array([np.quantile(R[b], 0.8) for b in bins])
 74    observed_slope = float(np.polyfit(kb, np.log(rb + 1e-30), 1)[0])
 75    predicted_slope = 2*np.log(rho)
 76    stability.append({"rho": rho, "predicted_log_slope": predicted_slope,
 77                      "observed_log_slope": observed_slope,
 78                      "relative_error": abs(observed_slope-predicted_slope)/abs(predicted_slope)})
 79
 80# Prediction 2: at rho=1 there is no exponential decay; rho>1 grows at rate 2 log rho.
 81boundary = []
 82for rho in [0.999, 1.0, 1.001, 1.01]:
 83    R = return_proxy(rho * np.exp(1j * omega[:1]), np.array([1.0]), ks)
 84    observed = float(np.polyfit(ks, np.log(R), 1)[0])
 85    predicted = 2*np.log(rho)
 86    boundary.append({"rho": rho, "predicted_log_slope": predicted, "observed_log_slope": observed,
 87                     "absolute_error": abs(observed-predicted)})
 88
 89# Prediction 3: increasing mode count improves quadrature of the pair-difference spectrum.
 90scaling = []
 91for M in [4, 8, 16, 32, 64, 128]:
 92    idx = select_frequency(M)
 93    R = return_proxy(teacher_lam[idx], q[idx], ks)
 94    scaling.append({"M": M, "log_rmse": relative_log_error(R_teacher, R),
 95                    "slope_relative_error": abs(fit_slope(ks, R)-fit_slope(ks, R_teacher))/max(abs(fit_slope(ks, R_teacher)),1e-12)})
 96
 97out = {"seed": SEED, "D": D, "ks": [1, 400], "teacher_slope": fit_slope(ks, R_teacher),
 98       "stability_sweep": stability, "boundary_sweep": boundary,
 99       "compression": compression, "mode_scaling": scaling}
100Path("results.json").write_text(json.dumps(out, indent=2))
101print(json.dumps(out, indent=2))