1import json
  2import time
  3from pathlib import Path
  4import numpy as np
  5
  6SEED = 3104
  7rng_global = np.random.default_rng(SEED)
  8
  9
 10def sample_rff(n, d=2, mode="gaussian", rng=None):
 11    rng = np.random.default_rng() if rng is None else rng
 12    if mode == "gaussian":
 13        # A conventional isotropic Gaussian RFF control with deliberately broad bandwidth.
 14        w = rng.normal(0.0, 3.0, size=(n, d))
 15    elif mode == "uniform":
 16        w = rng.uniform(-6.0, 6.0, size=(n, d))
 17    elif mode == "gevery_matched":
 18        # density proportional to exp(-2*kappa*||w||^(1/s)); s=1, kappa=.75.
 19        # In d dimensions, y=2*kappa*r^(1/s) is Gamma(d*s, 1).
 20        s, kappa = 1.0, 0.75
 21        y = rng.gamma(shape=d * s, scale=1.0, size=n)
 22        r = (y / (2.0 * kappa)) ** s
 23        direction = rng.normal(size=(n, d))
 24        direction /= np.linalg.norm(direction, axis=1, keepdims=True)
 25        w = r[:, None] * direction
 26    elif mode == "sobolev_matched":
 27        # density proportional to (1+||w||^2)^(-bar_s), bar_s>d/2.
 28        bar_s = 3.0
 29        u = rng.beta(d / 2.0, bar_s - d / 2.0, size=n)
 30        r = np.sqrt(u / (1.0 - u))
 31        direction = rng.normal(size=(n, d))
 32        direction /= np.linalg.norm(direction, axis=1, keepdims=True)
 33        w = r[:, None] * direction
 34    else:
 35        raise ValueError(mode)
 36    b = rng.uniform(0.0, 2.0 * np.pi, size=n)
 37    return w, b
 38
 39
 40def features(x, w, b):
 41    return np.sqrt(2.0) * np.cos(x @ w.T + b[None, :])
 42
 43
 44def target(x):
 45    # Smooth (indeed analytic) low-frequency periodic target, representative of a
 46    # coordinate field whose spectral energy is concentrated near the origin.
 47    return (0.80 * np.cos(x[:, 0] + 0.20) +
 48            0.55 * np.cos(x[:, 1] - 0.70) +
 49            0.35 * np.cos(2.0 * x[:, 0] + x[:, 1] + 0.40) +
 50            0.18 * np.cos(3.0 * x[:, 1] - 0.30))
 51
 52
 53def whiten(z, eps_scale=1e-5):
 54    g = (z.T @ z) / z.shape[0]
 55    eps = eps_scale * np.trace(g) / g.shape[0]
 56    vals, vecs = np.linalg.eigh(g + eps * np.eye(g.shape[0]))
 57    vals = np.maximum(vals, eps)
 58    zw = z @ (vecs * (1.0 / np.sqrt(vals))) @ vecs.T
 59    return zw, np.linalg.eigvalsh(g), float(eps)
 60
 61
 62def ridge_fit(z, y, alpha=1e-5):
 63    a = z.T @ z + alpha * np.eye(z.shape[1])
 64    return np.linalg.solve(a, z.T @ y)
 65
 66
 67def run():
 68    rng = np.random.default_rng(SEED)
 69    # Random train/test locations avoid exploiting a grid-specific aliasing artifact.
 70    xtr = rng.uniform(-np.pi, np.pi, size=(2048, 2))
 71    xte = rng.uniform(-np.pi, np.pi, size=(4096, 2))
 72    ytr, yte = target(xtr), target(xte)
 73    modes = ["gaussian", "uniform", "gevery_matched"]
 74    counts = [64, 128, 256, 512]
 75    rows = []
 76    t0 = time.perf_counter()
 77    # Three independent draws provide a small reproducibility check.
 78    for n in counts:
 79        for mode in modes:
 80            for whitened in [False, True]:
 81                errs, conds = [], []
 82                for rep in range(3):
 83                    rr = np.random.default_rng(SEED + 10000 * rep + n + hash(mode) % 997)
 84                    w, b = sample_rff(n, mode=mode, rng=rr)
 85                    ztr, zte = features(xtr, w, b), features(xte, w, b)
 86                    raw_eigs = np.linalg.eigvalsh((ztr.T @ ztr) / len(xtr))
 87                    cond = float(raw_eigs[-1] / max(raw_eigs[0], 1e-15))
 88                    if whitened:
 89                        ztr, _, _ = whiten(ztr)
 90                        # Whitening is an input-coordinate transform learned on train;
 91                        # use the same transform for test by recomputing explicitly below.
 92                        g = (features(xtr, w, b).T @ features(xtr, w, b)) / len(xtr)
 93                        eps = 1e-5 * np.trace(g) / n
 94                        vals, vecs = np.linalg.eigh(g + eps * np.eye(n))
 95                        zte = zte @ (vecs * (1.0 / np.sqrt(np.maximum(vals, eps)))) @ vecs.T
 96                    beta = ridge_fit(ztr, ytr)
 97                    pred = zte @ beta
 98                    errs.append(float(np.sqrt(np.mean((pred - yte) ** 2))))
 99                    conds.append(cond)
100                rows.append({"N": n, "mode": mode, "whitened": whitened,
101                             "test_rmse_mean": float(np.mean(errs)),
102                             "test_rmse_std": float(np.std(errs)),
103                             "raw_gram_condition_mean": float(np.mean(conds))})
104    # Cheap mathematical sanity checks: empirical tails should decrease in the
105    # expected order, and whitening should make the training Gram spectrum nearly flat.
106    tail_rng = np.random.default_rng(SEED + 77)
107    checks = {}
108    for mode in ["gaussian", "gevery_matched", "sobolev_matched"]:
109        w, _ = sample_rff(200000, mode=mode, rng=tail_rng)
110        r = np.linalg.norm(w, axis=1)
111        checks[mode] = {"median_radius": float(np.median(r)),
112                        "q99_radius": float(np.quantile(r, .99)),
113                        "mean_radius": float(np.mean(r))}
114    w, b = sample_rff(256, mode="gevery_matched", rng=np.random.default_rng(SEED))
115    zz = features(xtr, w, b)
116    before = np.linalg.eigvalsh((zz.T @ zz) / len(xtr))
117    zw, after, eps = whiten(zz)
118    after_reg = np.linalg.eigvalsh((zw.T @ zw) / len(xtr))
119    checks["whitening"] = {
120        "raw_condition": float(before[-1] / max(before[0], 1e-15)),
121        "whitened_condition": float(after_reg[-1] / max(after_reg[0], 1e-15)),
122        "whitened_eigen_min": float(after_reg[0]),
123        "whitened_eigen_max": float(after_reg[-1]),
124        "ridge_epsilon": eps,
125    }
126    out = {"seed": SEED, "runtime_sec": time.perf_counter() - t0,
127           "rows": rows, "sanity_checks": checks}
128    Path("results.json").write_text(json.dumps(out, indent=2))
129    print(json.dumps(out, indent=2))
130
131
132if __name__ == "__main__":
133    run()